diff --git a/docs/decisions/0019-python-context-compaction-strategy.md b/docs/decisions/0019-python-context-compaction-strategy.md
new file mode 100644
index 0000000000..11e1c091e5
--- /dev/null
+++ b/docs/decisions/0019-python-context-compaction-strategy.md
@@ -0,0 +1,1242 @@
+---
+status: accepted
+contact: eavanvalkenburg
+date: 2026-02-10
+deciders: eavanvalkenburg, markwallace-microsoft, sphenry, alliscode, johanst, brettcannon, westey-m
+consulted: taochenosu, moonbox3, dmytrostruk, giles17
+---
+
+# Context Compaction Strategy for Long-Running Agents
+
+## Context and Problem Statement
+
+Long-running agents need **context compaction** — automatically summarizing or truncating conversation history when approaching token limits. This is particularly important for agents that make many tool calls in succession (10s or 100s), where the context can grow unboundedly.
+
+[ADR-0016](0016-python-context-middleware.md) established the `ContextProvider` (hooks pattern) and `HistoryProvider` architecture for session management and context engineering. The .NET SDK comparison table notes:
+
+> **Message reduction**: `IChatReducer` on `InMemoryChatHistoryProvider` → Not yet designed (see Open Discussion: Context Compaction)
+
+This ADR proposes a design for context compaction that integrates with the chosen architecture.
+
+### Why Current Architecture Cannot Support In-Run Compaction
+
+An [analysis of the current message flow](https://gist.github.com/victordibia/ec3f3baf97345f7e47da025cf55b999f) identified three structural barriers to implementing compaction inside the tool loop:
+
+1. **History loaded once**: `HistoryProvider.get_messages()` is only called once during `before_run` at the start of `agent.run()`. The tool loop maintains its own message list internally and never re-reads from the provider.
+
+2. **`ChatMiddleware` modifies copies**: `ChatMiddleware` receives a **copy** of the message list each iteration. Clearing/replacing `context.messages` in middleware only affects that single LLM call — the tool loop's internal message list keeps growing with each tool result.
+
+3. **`FunctionMiddleware` wraps tool calls, not LLM calls**: `FunctionMiddleware` runs around individual tool executions, not around the LLM call that triggers them. It cannot modify the message history between iterations.
+
+```
+agent.run(task)
+ │
+ ├── ContextProvider.before_run() ← Load history, inject context ONCE
+ │
+ ├── chat_client.get_response(messages)
+ │ │
+ │ ├── messages = copy(messages) ← NEW list created
+ │ │
+ │ └── for attempt in range(max_iterations): ← TOOL LOOP
+ │ ├── ChatMiddleware(copy of messages) ← Modifies copy only
+ │ ├── LLM call(messages) ← Response may contain tool_calls
+ │ ├── FunctionMiddleware(tool_call) ← Wraps each tool execution
+ │ │ └── Execute single tool call
+ │ └── messages.extend(tool_results) ← List grows unbounded
+ │
+ └── ContextProvider.after_run() ← Store messages ONCE
+```
+
+**Consequence**: There is currently **no way** to compact messages during the tool loop such that subsequent LLM calls use the reduced context. Any middleware-based approach only affects individual LLM calls but the underlying list keeps growing.
+
+### Message-list correctness constraint: Atomic group preservation
+
+A critical correctness constraint for any compaction strategy: **tool calls and their results must be kept together**. LLM APIs (OpenAI, Azure, etc.) require that an assistant message containing `tool_calls` is always followed by corresponding `tool` result messages. A compaction strategy that removes one without the other will cause API errors. This is extended for reasoning models, at least in the OpenAI Responses API with a Reasoning content, without it you also get failed calls.
+
+Strategies must treat `[assistant message with tool_calls] + [tool result messages]` as atomic groups — either keep the entire group or remove it entirely. Option 1 addresses this structurally in both Variant C1 (precomputed `MessageGroups`) and Variant C2 (precomputed `_group_*` annotations on messages), so strategy authors do not need to rediscover raw boundaries on every pass.
+
+### Where Compaction Is Needed
+
+Compaction must be applicable in **three primary points** in the agent lifecycle:
+
+| Point | When | Purpose |
+|-------|------|---------|
+| **In-run** | During the (potentially) multiple calls to a ChatClient's `get_response` within a single `agent.run()` | Keep context within limits as tool calls accumulate and project only included messages per model call |
+| **Pre-write\*** | Before `HistoryProvider.save_messages()` in `after_run` | Compact before persisting to storage, limiting storage size, _only applies to messages from a run_ |
+| **On existing storage\*** | Outside of `agent.run()`, as a maintenance operation | Compact stored history (e.g., cron job, manual trigger) |
+
+**\***: Should pre-write and existing-storage compaction share one unified configuration/setup to reduce duplicate strategy wiring, and then either: each write overrides the full storage, or only new messages are compacted while a separate interface can be called to compact the existing storage?
+
+### Scope: Not Applicable to Service-Managed Storage
+
+**All compaction discussed in this ADR is irrelevant when using only service-managed storage** (`service_session_id` is set). In that scenario:
+- The service manages message history internally — the client never holds the full conversation
+- Only new messages are sent to/from the service each turn
+- The service is responsible for its own context window management and compaction
+- The client has no message list to compact
+
+This ADR applies to two scenarios where the **client** constructs and manages the message list sent to the model:
+
+1. **With local storage** (e.g., `InMemoryHistoryProvider`, Redis, Cosmos) — compaction is needed during a run, currently no compaction is done in our abstractions.
+2. **Without any storage** (`store=False`, no `HistoryProvider`) — in-run compaction is still critical for long-running, tool-heavy agent invocations where the message list grows unbounded within a single `agent.run()` call
+
+## Decision Drivers
+
+- **Applicable across primary points**: The strategy model must work at pre-write, in-run, and on existing storage, this means it must be:
+ - **Composable with HistoryProvider**: Works naturally with the `HistoryProvider` subclass from ADR-0016
+ - **Composable with function calling/chat clients**: Can be applied during the inner loop of the chat clients
+- **Message-list correctness**: Compaction must preserve required assistant/tool/result ordering and reasoning/tool-call pairings so the model input stays valid
+- **Chainable**/**Composable**: Multiple strategies must be composable (e.g., summarize older messages then truncate to fit token budget).
+
+## Considered Options
+
+- Standalone `CompactionStrategy` object composed into `HistoryProvider` and `ChatClient`
+- `CompactionStrategy` as a mixin for `HistoryProvider` subclasses
+- Separate `CompactionProvider` set directly on the agent
+- Mutable message access in `ChatMiddleware`
+
+
+## Pros and Cons of the Options
+
+### Option 1: Standalone `CompactionStrategy` Object
+
+Define an abstract `CompactionStrategy` that can be **composed into any `HistoryProvider`** and also passed to the agent for in-run compaction.
+
+There are three sub-variants for the method signature, which differ in mutability semantics and input structure, all of them use `__call__` to be easily used as a callable, and allow simple strategies to be expressed as simple functions, and if you need additional state or helper methods you can implement a class with `__call__`:
+
+#### Variant A: In-place mutation
+
+The strategy mutates the provided list directly and returns `bool` indicating whether compaction occurred. Zero-allocation in the no-op case, and the tool loop doesn't need to reassign the list.
+
+```python
+@runtime_checkable
+class CompactionStrategy(Protocol):
+ """Abstract strategy for compacting a list of messages in place."""
+
+ async def __call__(self, messages: list[Message]) -> bool:
+ """Compact messages in place. Returns True if compaction occurred."""
+ ...
+```
+
+#### Variant B: Return new list
+
+The strategy returns a new list (leaving the original unchanged) plus a `bool` indicating whether compaction occurred. This is safer when the caller needs the original list preserved (e.g., for logging or fallback), and is a more functional style that avoids side-effect surprises.
+
+```python
+@runtime_checkable
+class CompactionStrategy(Protocol):
+ """Abstract strategy for compacting a list of messages."""
+
+ async def __call__(self, messages: Sequence[Message]) -> tuple[list[Message], bool]:
+ """Return (compacted_messages, did_compact)."""
+ ...
+```
+
+Tool loop integration requires reassignment:
+
+```python
+# Inside the function invocation loop
+messages.append(tool_result_message)
+if compacter := config.get("compaction_strategy"):
+ compacted, did_compact = await compacter(messages)
+ if did_compact:
+ messages.clear()
+ messages.extend(compacted)
+```
+
+#### Variant C: Group-aware compaction entry points
+
+Variant C has two sub-variants that provide the same logical grouping behavior:
+- **C1 (`MessageGroups` state object):** group metadata lives in a sidecar container.
+- **C2 (`_`-prefixed message attributes):** group metadata lives directly on messages in `additional_properties`.
+
+Both approaches let strategies operate on logical units (`system`, `user`, `assistant_text`, `tool_call`) instead of re-deriving boundaries every time.
+
+##### Variant C1: `MessageGroups` sidecar state
+
+```python
+@dataclass
+class MessageGroup:
+ """A logical group of messages that must be kept or removed together."""
+ kind: Literal["system", "user", "assistant_text", "tool_call"]
+ messages: list[Message]
+
+ @property
+ def length(self) -> int:
+ """Number of messages in this group."""
+ return len(self.messages)
+
+
+@dataclass
+class MessageGroups:
+ groups: list[MessageGroup]
+
+ @classmethod
+ def from_messages(cls, messages: list[Message]) -> "MessageGroups":
+ """Build grouped state from a flat message list."""
+ groups: list[MessageGroup] = []
+ i = 0
+ while i < len(messages):
+ msg = messages[i]
+ if msg.role == "system":
+ groups.append(MessageGroup(kind="system", messages=[msg]))
+ i += 1
+ elif msg.role == "user":
+ groups.append(MessageGroup(kind="user", messages=[msg]))
+ i += 1
+ elif msg.role == "assistant" and getattr(msg, "tool_calls", None):
+ group_msgs = [msg]
+ i += 1
+ while i < len(messages) and messages[i].role == "tool":
+ group_msgs.append(messages[i])
+ i += 1
+ groups.append(MessageGroup(kind="tool_call", messages=group_msgs))
+ else:
+ groups.append(MessageGroup(kind="assistant_text", messages=[msg]))
+ i += 1
+ return cls(groups)
+
+ def summary(self) -> dict[str, int]:
+ return {
+ "group_count": len(self.groups),
+ "message_count": sum(len(g.messages) for g in self.groups),
+ "tool_call_count": sum(1 for g in self.groups if g.kind == "tool_call"),
+ }
+
+ def to_messages(self) -> list[Message]:
+ """Flatten grouped state back into a flat message list."""
+ return [msg for group in self.groups for msg in group.messages]
+
+
+class CompactionStrategy(Protocol):
+ """Callable strategy for group-aware compaction."""
+
+ async def __call__(self, groups: MessageGroups) -> bool:
+ """Compact by mutating grouped state. Returns True if changed.
+
+ Group kinds:
+ - "system": system message(s)
+ - "user": a single user message
+ - "assistant_text": an assistant message without tool calls
+ - "tool_call": an assistant message with tool_calls + all corresponding
+ tool result messages (atomic unit)
+ """
+ ...
+```
+
+Class-based strategies implement `__call__` directly:
+
+```python
+class ExcludeOldestGroupsStrategy:
+ async def __call__(self, groups: MessageGroups) -> bool:
+ # Mutate grouped state in place.
+ ...
+```
+
+The framework builds and flattens grouped state through `MessageGroups` methods:
+
+```python
+# Usage at a compaction point:
+groups = MessageGroups.from_messages(messages)
+logger.debug("Pre-compaction summary: %s", groups.summary())
+# optional also emit OTEL events next to these loggers, but not sure if needed
+await strategy(groups)
+logger.debug("Post-compaction summary: %s", groups.summary())
+response = await get_response(messages=groups.to_messages())
+# add messages from response into new group and to the groups.
+```
+
+**Note on in-run integration (C1):** Variant C1 requires maintaining grouped sidecar state (`MessageGroups` / underlying `list[MessageGroup]`) alongside the function-calling loop message list. Because `BaseChatClient` is stateless between calls, C1 cannot be cleanly implemented only in `BaseChatClient`; a stateful loop layer must own and update that grouped structure across roundtrips.
+
+##### Variant C2: `_`-prefixed metadata directly on `Message`
+
+Variant C2 achieves the same grouping behavior as C1 but stores grouping metadata on messages instead of in a sidecar `MessageGroups` object.
+
+```python
+def _annotate_groups(messages: list[Message]) -> None:
+ """Annotate messages with group metadata in additional_properties.
+
+ Metadata keys:
+ - "_group_id": stable group id for all messages in the same logical unit
+ - "_group_kind": "system" | "user" | "assistant_text" | "tool_call"
+ - "_group_index": order of groups in the current list
+ """
+ group_index = 0
+ i = 0
+ while i < len(messages):
+ msg = messages[i]
+ group_id = f"g-{group_index}"
+ if msg.role == "assistant" and getattr(msg, "tool_calls", None):
+ msg.additional_properties["_group_id"] = group_id
+ msg.additional_properties["_group_kind"] = "tool_call"
+ msg.additional_properties["_group_index"] = group_index
+ i += 1
+ while i < len(messages) and messages[i].role == "tool":
+ messages[i].additional_properties["_group_id"] = group_id
+ messages[i].additional_properties["_group_kind"] = "tool_call"
+ messages[i].additional_properties["_group_index"] = group_index
+ i += 1
+ else:
+ kind = (
+ "system" if msg.role == "system"
+ else "user" if msg.role == "user"
+ else "assistant_text"
+ )
+ msg.additional_properties["_group_id"] = group_id
+ msg.additional_properties["_group_kind"] = kind
+ msg.additional_properties["_group_index"] = group_index
+ i += 1
+ group_index += 1
+
+
+class CompactionStrategy(Protocol):
+ async def __call__(self, messages: list[Message]) -> bool:
+ """Compact using message annotations; mutate in place."""
+ ...
+```
+
+**Note on in-run integration (C2):** `BaseChatClient` should annotate new messages incrementally as they are appended (rather than re-running `_annotate_groups` over the full list every roundtrip). Unlike C1, C2 does not require a separate grouped sidecar in the function-calling loop; strategies can operate directly on `list[Message]` using `_group_*` metadata attached to the messages themselves. This makes C2 feasible as a fully `BaseChatClient`-localized implementation and provides a cleaner separation of responsibilities. In C2 and derived variants (D2/E2/F2), full ownership of compaction and message-attribute lifecycle belongs to the chat client to avoid double work: the chat client assigns/updates attributes (including `_group_id` for new tool-result messages added by function calling), and the function-calling layer remains unaware of this mechanism.
+
+#### Variant D: Exclude-based projection (builds on Variant C1/C2)
+
+Variant D also has two sub-variants:
+- **D1:** exclusion state on `MessageGroup`.
+- **D2:** exclusion state on message `_`-attributes.
+
+##### Variant D1: exclusion state on `MessageGroup`
+
+```python
+@dataclass
+class MessageGroup:
+ kind: Literal["system", "user", "assistant_text", "tool_call"]
+ messages: list[Message]
+ excluded: bool = False
+ exclude_reason: str | None = None
+
+
+@dataclass
+class MessageGroups:
+ groups: list[MessageGroup]
+
+ def summary(self) -> dict[str, int]:
+ return {
+ "group_count": len(self.groups),
+ "message_count": sum(len(g.messages) for g in self.groups),
+ "tool_call_count": sum(1 for g in self.groups if g.kind == "tool_call"),
+ "included_group_count": sum(1 for g in self.groups if not g.excluded),
+ "included_message_count": sum(len(g.messages) for g in self.groups if not g.excluded),
+ "included_tool_call_count": sum(
+ 1 for g in self.groups if g.kind == "tool_call" and not g.excluded
+ ),
+ }
+
+ def get_messages(self, *, excluded: bool = False) -> list[Message]:
+ if excluded:
+ return [msg for g in self.groups for msg in g.messages]
+ return [msg for g in self.groups if not g.excluded for msg in g.messages]
+
+ def included_messages(self) -> list[Message]:
+ return self.get_messages(excluded=False)
+```
+
+During compaction, strategies/orchestrators mutate `group.excluded`/`group.exclude_reason` (including re-including groups with `excluded=False`) instead of discarding data.
+
+##### Variant D2: exclusion state on message `_`-attributes
+
+```python
+def set_group_excluded(messages: list[Message], *, group_id: str, reason: str | None = None) -> None:
+ for msg in messages:
+ if msg.additional_properties.get("_group_id") == group_id:
+ msg.additional_properties["_excluded"] = True
+ msg.additional_properties["_exclude_reason"] = reason
+
+
+def clear_group_excluded(messages: list[Message], *, group_id: str) -> None:
+ for msg in messages:
+ if msg.additional_properties.get("_group_id") == group_id:
+ msg.additional_properties["_excluded"] = False
+ msg.additional_properties["_exclude_reason"] = None
+
+
+def included_messages(messages: list[Message]) -> list[Message]:
+ return [m for m in messages if not m.additional_properties.get("_excluded", False)]
+```
+
+In D2, strategies project included context by filtering on `_excluded` instead of filtering `MessageGroup` objects.
+
+#### Variant E: Tokenization and accounting (builds on Variant C1/C2)
+
+Variant E has two sub-variants:
+- **E1:** token rollups cached on `MessageGroup`/`MessageGroups`.
+- **E2:** token rollups cached directly on messages via `_`-attributes.
+
+##### Variant E1: token rollups on grouped state
+
+Variant E1 adds tokenization metadata and cached token rollups to grouped state. This is independent of exclusion: token-aware strategies can use token metrics even if no groups are excluded. When combined with Variant D, token budgets can be enforced against included messages.
+
+To make token-budget compaction deterministic:
+1. Before **every** `get_response` call in the tool loop, tokenize every message currently in `all_messages` (regardless of source).
+2. Persist per-content token counts in `content.additional_properties["_token_count"]`.
+3. Build/update grouped state from tokenized messages and use cached rollups for threshold checks and summaries.
+
+```python
+class TokenizerProtocol(Protocol):
+ def count_tokens(self, content: AIContent, *, model_id: str | None = None) -> int: ...
+
+
+@dataclass
+class MessageGroup:
+ kind: Literal["system", "user", "assistant_text", "tool_call"]
+ messages: list[Message]
+ _token_count_cache: int | None = None
+
+ def token_count(self) -> int:
+ if self._token_count_cache is None:
+ self._token_count_cache = sum(
+ content.additional_properties.get("_token_count", 0)
+ for message in self.messages
+ for content in message.contents
+ )
+ return self._token_count_cache
+
+
+@dataclass
+class MessageGroups:
+ groups: list[MessageGroup]
+ _total_tokens_cache: int | None = None
+
+ def total_tokens(self) -> int:
+ if self._total_tokens_cache is None:
+ self._total_tokens_cache = sum(group.token_count() for group in self.groups)
+ return self._total_tokens_cache
+
+ def summary(self) -> dict[str, int]:
+ return {
+ "group_count": len(self.groups),
+ "message_count": sum(len(g.messages) for g in self.groups),
+ "tool_call_count": sum(1 for g in self.groups if g.kind == "tool_call"),
+ "total_tokens": self.total_tokens(),
+ "tool_call_tokens": sum(g.token_count() for g in self.groups if g.kind == "tool_call"),
+ }
+```
+And the following helper method should also be added:
+
+```python
+def _to_tokenized_groups(
+ messages: list[Message], *, tokenizer: TokenizerProtocol
+) -> MessageGroups:
+ tokenize_messages(messages, tokenizer=tokenizer)
+ return MessageGroups.from_messages(messages)
+```
+
+##### Variant E2: token rollups on message `_`-attributes
+
+```python
+def annotate_token_counts(messages: list[Message], *, tokenizer: TokenizerProtocol) -> None:
+ for message in messages:
+ message_token_count = 0
+ for content in message.contents:
+ count = tokenizer.count_tokens(content)
+ content.additional_properties["_token_count"] = count
+ message_token_count += count
+ message.additional_properties["_message_token_count"] = message_token_count
+
+
+def sum_tokens_by_group(messages: list[Message]) -> dict[str, int]:
+ """Compute group totals on demand from `_message_token_count`."""
+ tokens_by_group: dict[str, int] = {}
+ for message in messages:
+ group_id = message.additional_properties["_group_id"]
+ tokens_by_group[group_id] = tokens_by_group.get(group_id, 0) + message.additional_properties.get(
+ "_message_token_count", 0
+ )
+ return tokens_by_group
+```
+
+In E2, strategies evaluate `_message_token_count`/`_token_count` directly from messages and compute per-group totals on demand via `_group_id` (instead of caching `_group_token_count` on every message). This avoids duplicated state and ambiguity when one copy is updated but others are stale. If needed for performance, the function-invocation loop can keep an ephemeral `dict[group_id, token_count]` alongside the annotated message list.
+
+#### Variant F: Combined projection + tokenization (C + D + E)
+
+Variant F has two sub-variants:
+- **F1:** combined model on `MessageGroups`.
+- **F2:** combined model on `_`-annotated messages.
+
+##### Variant F1: combined model on `MessageGroups`
+
+Variant F1 combines Variant C1's grouped interface, Variant D1's exclusion semantics, and Variant E1's token accounting in one integrated model. This gives one state container for projection (`excluded`) and budget control (`token_count`), while preserving full history for final-return and diagnostics.
+
+For Variant F1, `MessageGroups.from_messages(...)` accepts an optional tokenizer and handles both tokenization and grouping before strategy execution:
+
+```python
+class TokenizerProtocol(Protocol):
+ def count_tokens(self, content: AIContent, *, model_id: str | None = None) -> int: ...
+
+
+@dataclass
+class MessageGroup:
+ kind: Literal["system", "user", "assistant_text", "tool_call"]
+ messages: list[Message]
+ excluded: bool = False
+ exclude_reason: str | None = None
+ _token_count_cache: int | None = None
+
+ def token_count(self) -> int:
+ if self._token_count_cache is None:
+ self._token_count_cache = sum(
+ content.additional_properties.get("_token_count", 0)
+ for message in self.messages
+ for content in message.contents
+ )
+ return self._token_count_cache
+
+
+@dataclass
+class MessageGroups:
+ groups: list[MessageGroup]
+ _total_tokens_cache: int | None = None
+
+ @classmethod
+ def from_messages(
+ cls,
+ messages: list[Message],
+ *,
+ tokenizer: TokenizerProtocol | None = None,
+ ) -> "MessageGroups":
+ if tokenizer is not None:
+ tokenize_messages(messages, tokenizer=tokenizer)
+ groups: list[MessageGroup] = []
+ i = 0
+ while i < len(messages):
+ msg = messages[i]
+ if msg.role == "system":
+ groups.append(MessageGroup(kind="system", messages=[msg]))
+ i += 1
+ elif msg.role == "user":
+ groups.append(MessageGroup(kind="user", messages=[msg]))
+ i += 1
+ elif msg.role == "assistant" and getattr(msg, "tool_calls", None):
+ group_msgs = [msg]
+ i += 1
+ while i < len(messages) and messages[i].role == "tool":
+ group_msgs.append(messages[i])
+ i += 1
+ groups.append(MessageGroup(kind="tool_call", messages=group_msgs))
+ else:
+ groups.append(MessageGroup(kind="assistant_text", messages=[msg]))
+ i += 1
+ return cls(groups)
+
+ def get_messages(self, *, excluded: bool = False) -> list[Message]:
+ if excluded:
+ return [msg for g in self.groups for msg in g.messages]
+ return [msg for g in self.groups if not g.excluded for msg in g.messages]
+
+ def included_messages(self) -> list[Message]:
+ return self.get_messages(excluded=False)
+
+ def total_tokens(self) -> int:
+ if self._total_tokens_cache is None:
+ self._total_tokens_cache = sum(group.token_count() for group in self.groups)
+ return self._total_tokens_cache
+
+ def included_token_count(self) -> int:
+ return sum(g.token_count() for g in self.groups if not g.excluded)
+
+ def summary(self) -> dict[str, int]:
+ return {
+ "group_count": len(self.groups),
+ "message_count": sum(len(g.messages) for g in self.groups),
+ "tool_call_count": sum(1 for g in self.groups if g.kind == "tool_call"),
+ "included_group_count": sum(1 for g in self.groups if not g.excluded),
+ "included_message_count": sum(len(g.messages) for g in self.groups if not g.excluded),
+ "included_tool_call_count": sum(
+ 1 for g in self.groups if g.kind == "tool_call" and not g.excluded
+ ),
+ "total_tokens": self.total_tokens(),
+ "tool_call_tokens": sum(g.token_count() for g in self.groups if g.kind == "tool_call"),
+ "included_tokens": self.included_token_count(),
+ }
+
+
+class CompactionStrategy(Protocol):
+ async def __call__(self, groups: MessageGroups) -> None:
+ """Mutate the provided groups in place."""
+ ...
+```
+
+##### Variant F2: combined model on `_`-annotated messages
+
+```python
+class CompactionStrategy(Protocol):
+ async def __call__(self, messages: list[Message]) -> bool:
+ """Mutate message annotations in place."""
+ ...
+
+
+async def compact_with_annotations(
+ messages: list[Message], *, strategy: CompactionStrategy, tokenizer: TokenizerProtocol
+) -> list[Message]:
+ # C2: annotate group boundaries
+ _annotate_groups(messages)
+ # E2: annotate token metrics
+ annotate_token_counts(messages, tokenizer=tokenizer)
+ _ = sum_tokens_by_group(messages) # optional ephemeral aggregate in loop state
+
+ # D2/F2: strategy toggles _excluded/_exclude_reason and can rewrite messages
+ _ = await strategy(messages)
+
+ # Project only included messages for model call
+ return [m for m in messages if not m.additional_properties.get("_excluded", False)]
+```
+
+F2 avoids a sidecar object but requires strict ownership rules for `_` attributes (who sets, updates, clears, and validates them). To prevent duplicate work and drift, this ownership should live entirely in `BaseChatClient`, while the function-calling layer remains attribute-unaware.
+
+**Trade-offs between variants:**
+
+| Aspect | Variant A (in-place) | Variant B (return new) | Variant C1 (`MessageGroups`) | Variant C2 (`_` attrs) | Variant D1 (`MessageGroups` exclude) | Variant D2 (`_excluded` attrs) | Variant E1 (group token caches) | Variant E2 (message token attrs + on-demand group sums) | Variant F1 (`MessageGroups` combined) | Variant F2 (`_` attrs combined) |
+|--------|---------------------|----------------------|-------------------------------|-----------------------|--------------------------------------|-------------------------------|----------------------------------|-------------------------------------|-----------------------------------|----------------------------------|
+| **Allocation** | Zero in no-op case | Always allocates tuple | Grouping sidecar allocation | No sidecar; metadata writes | D1 + exclusion state | D2 + metadata writes | E1 + token cache sidecar | E2 + message metadata writes | Highest sidecar state | No sidecar; highest metadata writes |
+| **Safety** | Caller loses original | Original preserved | State isolated in sidecar | Metadata mutates source messages | Full grouped history preserved | Full message history preserved | Deterministic token rollups in sidecar | Deterministic token rollups on messages | Strong isolation of all compaction state | Shared-message mutation can leak across layers |
+| **Strategy complexity** | Must handle atomic groups | Must handle atomic groups | Groups pre-computed by framework | Reads `_group_*` fields | Exclude/re-include by group | Exclude/re-include by `_group_id` | Token budget via group APIs | Token budget via `_token*` fields | Unified exclude + token policy via group APIs | Unified policy via many message attrs |
+| **Chaining** | Natural (same list) | Pipe output to next input | Natural (same group state) | Natural (same annotated message list) | Natural | Natural | Natural | Natural | Natural | Natural |
+| **Framework complexity** | Minimal | Reassignment logic | Grouping + flattening layer | Annotation lifecycle/validation | C1 + exclusion semantics | C2 + projection/filter semantics | C1 + tokenizer + cache invalidation | C2 + tokenizer + attr invalidation | Highest sidecar orchestration | Highest attr lifecycle orchestration |
+
+**Usage with `HistoryProvider`:**
+
+The `compaction_strategy` parameter accepts either a single `CompactionStrategy` or it can take a composed/chained strategy.
+
+```python
+
+class HistoryProvider(ContextProvider):
+ def __init__(
+ self,
+ source_id: str,
+ *,
+ load_messages: bool = True,
+ store_inputs: bool = True,
+ store_responses: bool = True,
+ store_excluded_messages: bool = True, # NEW: persist excluded groups/messages or only included
+ # NEW: optional compaction strategy, can be a single strategy or a chained/composed strategy
+ compaction_strategy: CompactionStrategy | None = None,
+ # NEW: optional tokenizer for token-aware compaction strategies
+ tokenizer: TokenizerProtocol | None = None,
+ ): ...
+
+ async def after_run(self, agent, session, context, state) -> None:
+ messages_to_store = self._collect_messages(context)
+ groups = MessageGroups.from_messages(messages_to_store, tokenizer=self.tokenizer)
+ if self.compaction_strategy:
+ await self.compaction_strategy(groups)
+ messages_to_store = groups.get_messages(excluded=self.store_excluded_messages)
+ if messages_to_store:
+ await self.save_messages(context.session_id, messages_to_store)
+```
+
+**Simple usage:**
+
+```python
+strategy = SlidingWindowStrategy(max_messages=100)
+
+agent = client.create_agent(
+ context_providers=[
+ InMemoryHistoryProvider("memory", compaction_strategy=strategy),
+ ],
+)
+```
+
+There are two ways we can do this:
+1. Before writing to storage in `after_run`, compaction is called on the new messages,
+ combined with: a new `compact` method, that reads the full history, calls the compaction strategy with the full history, then writes the compacted result back to storage (also requires a `overwrite` flag on the `save_messages` method). This makes removing old messages from storage a explicit action that the user initiaties instead of being implicitly triggered by `after_run` writes, but it also means compaction strategies only see new messages instead of the full history (unless they read it themselves), the `compact` method could then also have a override for the strategy to use (and/or the tokenizer in case of Variant E1/E2/F1/F2).
+
+ ```python
+ class HistoryProvider(ContextProvider):
+ ...
+ async def compact(self, session_id: str, *, strategy: CompactionStrategy | None = None, tokenizer: TokenizerProtocol | None = None) -> None:
+ history = await self.get_messages(session_id)
+ if tokenizer:
+ tokenize_messages(history, tokenizer=tokenizer)
+ applicable_strategy = strategy or self.compaction_strategy
+ await applicable_strategy(history) # compaction mutates history in place or returns new list depending on variant
+ await self.save_messages(session_id, history, overwrite=True) # write compacted history back to storage
+ ```
+
+2. Before writing the history is loaded (could already be in-memory from `before_run`), compaction is called on the full history (old + new), then the compacted result is written back to storage. This allows compaction strategies to consider the full history when deciding what to keep, but it also means the provider needs to support writing the full history back (not just appending new messages).
+
+Given the explicit nature, and the ability to do the heavy lifting of reading, compacting and writing outside of the agent loop, we decide to go with the first setup, if we decide to use Option 1 overall.
+
+**Usage for in-run compaction (BaseChatClient):**
+
+In-run compaction should execute in `BaseChatClient` before every `get_response` call, regardless of whether function calling is enabled. This makes compaction behavior uniform for single-shot and looped invocations.
+
+For token-aware variants (E1/E2/F1/F2), a tokenizer must be configured because token counts are part of compaction decisions. For the grouped-state path (F1), use `MessageGroups.from_messages(..., tokenizer=...)` so tokenization and grouping happen together before strategy invocation.
+
+For C2/D2/E2/F2 specifically, `BaseChatClient` is the sole owner of compaction + `_`-attribute lifecycle. It should assume this work is required, annotate/refresh metadata on appended messages (including tool-result messages coming from function calling), and project included messages for model calls. The function-calling layer should not implement or duplicate any part of this mechanism.
+
+```python
+class BaseChatClient:
+ # NEW attributes on the existing class
+ compaction_strategy: CompactionStrategy | None = None
+ tokenizer: TokenizerProtocol | None = None # required for token-aware variants
+```
+
+Agent attributes stay the same and are passed into the chat client (similar to `ChatMiddleware` propagation):
+
+```python
+agent = Agent(
+ client=chat_client,
+ context_providers=[
+ InMemoryHistoryProvider("memory", compaction_strategy=boundary_strategy),
+ ],
+ compaction_strategy=compaction_strategy,
+ tokenizer=model_tokenizer, # required for token-aware variants (E1/E2/F1/F2)
+)
+
+chat_client.compaction_strategy = agent.compaction_strategy
+chat_client.tokenizer = agent.tokenizer
+```
+
+Execution then lives in `BaseChatClient.get_response(...)`:
+
+```python
+def get_response(
+ self,
+ messages: Sequence[Message],
+ *,
+ stream: bool = False,
+ options: Mapping[str, Any] | None = None,
+ **kwargs: Any,
+) -> Awaitable[ChatResponse[Any]] | ResponseStream[ChatResponseUpdate, ChatResponse[Any]]:
+ if not self.compaction_strategy:
+ return self._inner_get_response(
+ messages=messages,
+ stream=stream,
+ options=options or {},
+ **kwargs,
+ )
+
+ groups = MessageGroups.from_messages(
+ messages,
+ tokenizer=self.tokenizer,
+ )
+ # Compaction hook runs here and updates included/excluded state on groups.
+ projected = groups.included_messages()
+ return self._inner_get_response(
+ messages=projected,
+ stream=stream,
+ options=options or {},
+ **kwargs,
+ )
+```
+
+`BaseChatClient` always keeps the full grouped state (included + excluded) in memory and uses only the projected included messages for model calls. Return/persistence policy is handled outside the client (e.g., `HistoryProvider.store_excluded_messages`).
+
+When function calling is enabled, every model roundtrip still goes through `BaseChatClient.get_response(...)`, so compaction runs automatically without duplicating logic in function-invocation code.
+
+**Built-in strategies:**
+
+```python
+class TruncationStrategy(CompactionStrategy):
+ """Keep the last N messages, optionally preserving the system message."""
+ def __init__(self, *, max_messages: int, max_tokens: int, preserve_system: bool = True): ...
+
+class SlidingWindowStrategy(CompactionStrategy):
+ """Keep system message + last N messages."""
+ def __init__(self, *, max_messages: int, max_tokens: int): ...
+
+class SummarizationStrategy(CompactionStrategy):
+ """Summarize older messages using an LLM."""
+ def __init__(self, *, client: ..., max_messages_before_summary: int, max_tokens_before_summary: int): ...
+
+# etc
+```
+
+**Opinionated token budget based composed strategy pattern (Variant F1/F2):**
+
+This ADR proposes shipping a built-in composed strategy that enforces a token budget by running a list of regular strategies from top to bottom until the conversation fits the budget. This is intentionally opinionated and serves as a practical default/inspiration; advanced users can still implement custom orchestration logic. In F1, this strategy should drive `MessageGroup.excluded`; in F2, it should drive message `_excluded` annotations so model calls project only included context while preserving the full list.
+
+```python
+class TokenBudgetComposedStrategy(CompactionStrategy):
+ def __init__(
+ self,
+ *,
+ token_budget: int,
+ strategies: Sequence[CompactionStrategy],
+ early_stop: bool = False, # optional flag to stop after first strategy that meets the budget, or run all strategies regardless
+ ):
+ self.token_budget = token_budget
+ self.strategies = strategies
+ self.early_stop = early_stop
+
+ async def __call__(self, groups: MessageGroups) -> None:
+ if groups.included_token_count() <= self.token_budget:
+ return
+
+ for strategy in self.strategies:
+ await strategy(groups)
+
+ if self.early_stop and groups.included_token_count() <= self.token_budget:
+ break
+```
+
+This pattern keeps composition explicit and deterministic: ordered strategies, shared token metric, exclusion-flag semantics, optional re-inclusion by later strategies, and early stop as soon as budget is satisfied.
+
+- Good, because the same strategy model works at the three primary compaction points (pre-write, in-run, existing storage)
+- Good, because strategies are fully reusable — one instance can be shared across providers and agents
+- Good, because new strategies can be added without modifying `HistoryProvider`
+- Good, because with Variant A (in-place), the tool loop integration is zero-allocation in the no-op case
+- Good, because with Variant B (return new list), the caller retains the original list for logging or fallback
+- Good, because with Variants C1-F1 (grouped-state), strategy authors don't need to implement atomic group preservation — the framework handles grouping/flattening, making strategies simpler and less error-prone
+- Good, because with Variants C2-F2 (message annotations), we can avoid a sidecar `MessageGroups` container while still preserving logical groups through `_group_*` attributes
+- Good, because it is easy to test strategies in isolation
+- Good, because strategies can inspect `source_id` attribution on messages for informed decisions
+- Good, because in-run settings can be first-class `Agent` parameters and are propagated into `BaseChatClient` attributes
+- Good, because **chaining is natural** — for Variants A/C1-F2, each strategy mutates the same shared state in sequence; for Variant B, output pipes into the next input
+- Neutral, because Variants C1-F2 add framework complexity (grouping/flattening or annotation lifecycle, plus tokenization/exclusion accounting) but reduce strategy complexity
+- Bad, because it adds a new concept (`CompactionStrategy`) alongside the existing `ContextProvider`/`HistoryProvider` hierarchy
+- Bad, because Variants C1-F1 introduce a `MessageGroup` model that must stay in sync with any future message role changes
+- Bad, because Variants C2-F2 depend on careful `_`-attribute lifecycle management to avoid stale or inconsistent annotations
+
+### Option 2: `CompactionStrategy` as a Mixin for `HistoryProvider`
+
+Define compaction behavior as a mixin that `HistoryProvider` subclasses can opt into. The mixin adds `compact()` as an overridable method.
+
+```python
+class CompactingHistoryMixin:
+ """Mixin that adds compaction to a HistoryProvider."""
+
+ async def compact(self, messages: Sequence[ChatMessage]) -> list[ChatMessage]:
+ """Override to implement compaction logic. Default: no-op."""
+ return list(messages)
+
+
+class InMemoryHistoryProvider(CompactingHistoryMixin, HistoryProvider):
+ """In-memory history with compaction support."""
+
+ def __init__(
+ self,
+ source_id: str,
+ *,
+ max_messages: int | None = None,
+ **kwargs,
+ ):
+ super().__init__(source_id, **kwargs)
+ self.max_messages = max_messages
+
+ async def compact(self, messages: Sequence[ChatMessage]) -> list[ChatMessage]:
+ if self.max_messages and len(messages) > self.max_messages:
+ return list(messages[-self.max_messages:])
+ return list(messages)
+```
+
+The base `HistoryProvider` checks for the mixin and calls `compact()` at the right points:
+
+```python
+class HistoryProvider(ContextProvider):
+ async def before_run(self, agent, session, context, state) -> None:
+ history = await self.get_messages(context.session_id)
+ if isinstance(self, CompactingHistoryMixin):
+ history = await self.compact(history)
+ context.extend_messages(self.source_id, history)
+```
+
+For in-run compaction, `BaseChatClient` attributes would reference the provider's `compact()` method, but this requires knowing which provider to use:
+
+```python
+# Awkward: must extract compaction from a specific provider
+compacting_provider = next(
+ (p for p in agent._context_providers if isinstance(p, CompactingHistoryMixin)),
+ None,
+)
+base_chat_client.compaction_strategy = compacting_provider # provider IS the strategy
+```
+
+For existing storage:
+
+```python
+# Provider must implement CompactingHistoryMixin
+provider = InMemoryHistoryProvider("memory", max_messages=100)
+history = await provider.get_messages(session_id)
+compacted = await provider.compact(history)
+await provider.save_messages(session_id, compacted)
+```
+
+- Good, because no new top-level concept — compaction is part of the provider
+- Good, because the provider controls its own compaction logic
+- Neutral, because mixins are idiomatic Python but can be harder to reason about in complex hierarchies
+- Bad, because **compaction strategy is coupled to the provider** — cannot share the same strategy across different providers, or in-run.
+- Bad, because different strategies per compaction point (pre-write vs existing) require additional configuration or separate methods
+- Bad, because in-run compaction via `BaseChatClient` attributes requires extracting the mixin from the provider list — unclear which one to use if multiple exist
+- Bad, because `isinstance` checks are fragile and don't compose well
+- Bad, because testing compaction requires instantiating a full provider rather than testing the strategy in isolation
+- Bad, because existing storage compaction requires having the right provider type, not just any strategy
+- Bad, because **chaining is difficult** — compaction logic is embedded in the provider's `compact()` override, so composing multiple strategies (e.g., summarize then truncate) requires subclass nesting or manual delegation within a single `compact()` method, rather than declarative composition
+
+### Option 3: Separate `CompactionProvider` Set on the Agent
+
+Define compaction as a special `ContextProvider` subclass that the agent calls at all compaction points (pre-load, pre-write, in-run (calls `compact`), existing storage). It is added to the agent's `context_providers` list like any other provider.
+
+```python
+class CompactionProvider(ContextProvider):
+ """Context provider specialized for compaction.
+
+ Unlike regular ContextProviders, CompactionProvider is also invoked
+ during the function calling loop and can be used for storage maintenance.
+ """
+
+ @abstractmethod
+ async def compact(self, messages: Sequence[ChatMessage]) -> list[ChatMessage]:
+ """Reduce a list of messages."""
+ ...
+
+ async def before_run(self, agent, session, context, state) -> None:
+ """Compact messages loaded by previous providers before model invocation."""
+ all_messages = context.get_all_messages()
+ compacted = await self.compact(all_messages)
+ context.replace_messages(compacted)
+
+ async def after_run(self, agent, session, context, state) -> None:
+ """No-op by default. Subclasses can override for pre-write behavior."""
+ pass
+```
+
+**Usage:**
+
+```python
+agent = ChatAgent(
+ chat_client=client,
+ context_providers=[
+ InMemoryHistoryProvider("memory"), # Loads history
+ RAGContextProvider("rag"), # Adds RAG context
+ SlidingWindowCompaction("compaction", max_messages=100), # Compacts everything
+ ],
+)
+```
+
+The agent recognizes `CompactionProvider` instances and wires `compact()` into `BaseChatClient` attributes:
+
+```python
+class ChatAgent:
+ def _configure_base_chat_client(self, base_client: BaseChatClient) -> None:
+ compactors = [p for p in self._context_providers if isinstance(p, CompactionProvider)]
+ strategy = compactors[0] if compactors else None # Which one if multiple?
+ base_client.compaction_strategy = strategy
+```
+
+For existing storage, the `compact()` method is called directly:
+
+```python
+compactor = SlidingWindowCompaction("compaction", max_messages=100)
+history = await my_history_provider.get_messages(session_id)
+compacted = await compactor.compact(history)
+await my_history_provider.save_messages(session_id, compacted)
+```
+
+- Good, because it lives within the existing `ContextProvider` pipeline — no new concept
+- Good, because ordering relative to other providers is explicit (runs after RAG provider, etc.)
+- Good, because `before_run` can compact the combined output of all prior providers (history + RAG)
+- Good, because the `compact()` method works standalone for existing storage maintenance
+- Neutral, because **chaining is partially supported** — multiple `CompactionProvider` instances can be added to the provider list and will run in order during `before_run`/`after_run`, but in-run compaction via `BaseChatClient` attributes only wires a single strategy (which one to pick is ambiguous), so chaining works at boundaries but not during the tool loop
+- Bad, because the `CompactionProvider` has **dual roles** (context provider + compaction strategy), which muddies the ContextProvider contract
+- Bad, because `context.replace_messages()` is a new operation that doesn't exist today and conflicts with the append-only design of `SessionContext`
+- Bad, because in-run compaction still requires `isinstance` checks to wire into `BaseChatClient` attributes
+- Bad, because ordering sensitivity is subtle — must come after storage providers but before model invocation
+- Bad, because a `CompactionProvider` as a context provider gets `before_run`/`after_run` calls even when only its `compact()` method is needed (in-run and storage maintenance)
+
+### Option 4: Mutable Message Access in `ChatMiddleware`
+
+Instead of introducing a new compaction abstraction, change `ChatMiddleware` so that it can **replace the actual message list** used by the tool loop, rather than modifying a copy. This makes the existing middleware pattern sufficient for in-run compaction.
+
+**Required changes to the tool loop:**
+
+```python
+# Inside the function invocation loop
+# Current: ChatMiddleware modifies a copy, tool loop keeps its own list
+# Proposed: ChatMiddleware can replace the list, tool loop uses the replacement
+
+for attempt_idx in range(max_iterations):
+ context = ChatContext(messages=messages)
+ response = await middleware_pipeline.process(context)
+
+ # NEW: if middleware replaced messages, use the replacement
+ messages = context.messages # May be a new, compacted list
+
+ messages.extend(tool_results)
+```
+
+**Usage:**
+
+```python
+@chat_middleware
+async def compacting_middleware(context: ChatContext, next):
+ if count_tokens(context.messages) > budget:
+ compacted = compact(context.messages)
+ context.messages.clear()
+ context.messages.extend(compacted) # Persists because tool loop reads back
+ await next(context)
+
+agent = chat_client.create_agent(
+ middleware=[compacting_middleware],
+)
+```
+
+For boundary compaction, the same middleware runs at the chat client level. For existing storage compaction, a standalone utility function is needed since middleware only runs during `agent.run()`.
+
+- Good, because it uses the **existing `ChatMiddleware` pattern** — no new compaction concept
+- Good, because middleware already runs between LLM calls in the tool loop — it just needs the mutations to stick
+- Good, because users familiar with middleware get compaction "for free"
+- Neutral, because **chaining is implicit** — multiple compaction middleware can be stacked and will run in pipeline order, but there is no explicit composition model; middleware interact through side effects (mutating the shared message list) rather than declarative input/output, making chain behavior harder to reason about and debug
+- Bad, because it requires **changing how the tool loop manages messages** — the current copy-based architecture must be rethought
+- Bad, because multiple middleware could conflict when replacing messages (no coordination)
+- Bad, because it does **not cover existing storage compaction**
+- Bad, because it does **not cover pre-write compaction** — `ChatMiddleware` runs before the LLM call, not after `ContextProvider.after_run()`
+- Bad, because message replacement semantics in middleware are implicit (mutating a list) rather than explicit (returning a new list)
+- Bad, because it requires significant internal refactoring of the copy-based message flow in the function invocation layer
+
+
+## Decision Outcome
+
+Chosen option: **Option 1: Standalone `CompactionStrategy` Object** with **F2** (`_`-annotated messages) as the primary implementation model. We still document F1 as a valid alternative, but F2 is preferred because it introduces one less concept (no sidecar `MessageGroups` container), aligns with `BaseChatClient` statelessness by carrying state on messages themselves, and allows in-run compaction to stay localized to `BaseChatClient` rather than requiring extra grouped-state ownership in the function-calling loop.
+
+## Comparison to .NET Implementation
+
+The .NET SDK uses `IChatReducer` composed into `InMemoryChatHistoryProvider`:
+
+| Aspect | .NET | Proposed Options |
+|--------|------|-----------------|
+| Interface | `IChatReducer` with `ReduceAsync(messages) -> messages` | `CompactionStrategy.compact()` with three signature variants (Options 1-3) / `ChatMiddleware` mutation (Option 4) |
+| Attachment | Property on `InMemoryChatHistoryProvider` | Composed into `HistoryProvider` (Option 1) / mixin (Option 2) / separate provider (Option 3) / middleware (Option 4) |
+| Trigger | `ChatReducerTriggerEvent` enum: `AfterMessageAdded`, `BeforeMessagesRetrieval` | Pre-write + in-run + storage maintenance (Options 1-3 primary scope); post-load-style behavior can be covered by in-run pre-send projection |
+| Scope | Only within `InMemoryChatHistoryProvider` | Applicable to any `HistoryProvider` and the tool loop (Option 1) |
+
+Option 1's `CompactionStrategy` is the closest equivalent to .NET's `IChatReducer`, with a broader scope.
+
+### Achieving the same scenarios in MEAI/.NET
+
+| Python scenario | .NET/MEAI mechanism | How it maps |
+|-----------------|---------------------|-------------|
+| **Pre-write compaction** | `InMemoryChatHistoryProvider` + `ChatReducerTriggerEvent.AfterMessageAdded` | Reducer runs in `StoreChatHistoryAsync` after new request/response messages are added to storage (closest equivalent to pre-write persistence compaction). |
+| **Agent-level whole-list compaction (pre-send overlap with post-load)** | `ChatClientAgent` message assembly + chat-client decoration via `clientFactory` / `ChatClientAgentRunOptions.ChatClientFactory` | `ChatClientAgent` builds the full invocation message list (`ChatHistoryProvider` + `AIContextProviders` + input). A delegating `IChatClient` can compact that assembled list immediately before forwarding `GetResponseAsync`. |
+| **In-run compaction before every `get_response` call** | Base chat-client layer + delegating `IChatClient` wrapper | Compaction is executed in the base chat client before every `GetResponseAsync` call, so both single-shot and function-calling roundtrips get the same behavior. |
+| **Variant C1 grouped-state maintenance (`MessageGroup`)** | Keep grouped state in the same function-invocation/delegating-chat-client layer | Maintain and update grouped state across loop iterations in that layer, then flatten only for model calls. |
+| **Variant C2 message-annotation maintenance (`_group_*`)** | Keep message annotations in the same function-invocation/delegating-chat-client layer | Incrementally annotate newly appended messages with `_group_id`, `_group_kind`, and related metadata; filter/project directly from annotated message lists. |
+| **Compaction on existing storage** | `InMemoryChatHistoryProvider.GetMessages(...)` + `SetMessages(...)` (or custom provider equivalent) | Read stored history, apply reducer/strategy, and write back compacted history as a maintenance operation. |
+
+### Coverage Matrix
+
+How each option addresses the three primary compaction points and the current architectural limitations:
+
+| Compaction Point | Option 1 (Strategy) | Option 2 (Mixin) | Option 3 (Provider) | Option 4 (Middleware) |
+|-----------------|---------------------|-------------------|---------------------|-----------------------|
+| **Pre-write** | ✅ `HistoryProvider` param | ⚠️ Needs extra method | ⚠️ `after_run` override | ❌ Not supported |
+| **In-run (tool loop)** | ✅ `BaseChatClient` attrs | ⚠️ Awkward extraction | ⚠️ `isinstance` wiring | ⚠️ Requires refactoring copy semantics |
+| **Existing storage** | ✅ Standalone `compact()` | ✅ Provider's `compact()` | ✅ Standalone `compact()` | ❌ Not supported |
+| **Solves copy problem** | ✅ Runs inside loop | ⚠️ Indirectly | ⚠️ Indirectly | ⚠️ Requires deep refactor |
+| **Chaining** | ✅ Natural composition via wrapper | ❌ Coupled to provider | ⚠️ Boundary only, not in-run | ⚠️ Implicit via stacking |
+| **New concepts** | 1 (`CompactionStrategy`) | 1 (mixin) | 0.5 (reuses `ContextProvider`, but adds new method) | 0 (reuses `ChatMiddleware`) |
+
+
+## Appendix
+
+### Appendix A: Strategy and constraint background
+
+### Compaction Strategies (Examples)
+
+A compaction strategy takes a list of messages and returns a (potentially shorter) list, in almost all cases, there is certain logic that needs to be applied universally, such as retaining system messages, not breaking up function call and result pairs (for Responses that includes Reasoning as well, see [context section above](#message-list-correctness-constraint-atomic-group-preservation) for more info) as tool calls, etc. Beyond that, strategies can be as simple or complex as needed:
+
+- **Truncation**: Keep only the last N messages or N tokens, this is a likely done as a kind of zigzag, where the history grows, then get's truncated to some value below the token limit, then grows again, etc. This can be done on a simple message count basis, a character count basis, or more complex token counting basis.
+- **Summarization**: Replace older messages with an LLM-generated summary (depending on the implementation this could be done, by replacing the summarized messages, or by inserting a summary message in between and not loading messages older then the summarized ones)
+- **Selective removal**: Remove tool call/result pairs while keeping user/assistant turns
+- **Sliding window with anchor**: Keep system message + last N messages
+- **Custom logic**: The design should be extendible so that users can implement their own strategies.
+
+### Leveraging Source Attribution
+
+[ADR-0016](./0016-python-context-middleware.md#4-source-attribution-via-source_id) introduces `source_id` attribution on messages — each message tracks which `ContextProvider` added it. Compaction strategies can use this attribution to make informed decisions about what to compact and what to preserve:
+
+- **Preserve RAG context**: Messages from a RAG provider (e.g. `source_id: "rag"`) may be critical and should survive compaction
+- **Remove ephemeral context**: Messages marked as ephemeral (e.g., `source_id: "time"`) can be safely removed
+- **Protect user input**: Messages without a `source_id` (direct user input) should typically be preserved
+- **Selective tool result compaction**: Tool results from specific providers can be summarized while others are kept verbatim
+
+This means strategies don't need to rely solely on message position or role — they can make semantically meaningful compaction decisions based on the origin of each message.
+
+### Appendix B: Additional implementation notes
+
+#### Trigger mechanism for in-run compaction
+
+Running compaction after **every** tool call is wasteful — most iterations the context is well within limits. Instead, compaction should only trigger when a threshold is exceeded. There are several approaches to consider:
+
+1. **Message count threshold**: Trigger when the message list exceeds N messages. Simple to implement and predictable, but message count is a poor proxy for token usage — a single tool result can contain thousands of tokens while counting as one message.
+
+2. **Character/token count threshold**: Trigger when the estimated token count exceeds a budget. More accurate but requires a token counting mechanism (exact tokenization is model-specific and expensive; character-based heuristics like `len(text) / 4` are fast but approximate).
+
+3. **Iteration-based**: Trigger every N tool loop iterations (e.g., every 10th iteration). Predictable cadence but doesn't account for actual context growth — 10 iterations with small results may not need compaction while 3 iterations with large results might.
+
+4. **Strategy-internal**: Let the `CompactionStrategy.compact()` method decide internally — it receives the full message list and can return it unchanged if no compaction is needed. This is the simplest integration point (always call `compact()`, let the strategy no-op when appropriate) but has the overhead of calling into the strategy every iteration.
+
+The recommended approach is **strategy-internal with a lightweight guard**: the `compact()` method is called after each tool result, but strategy implementations should include a fast short-circuit check (e.g., `if len(messages) < self.threshold: return False`) to minimize overhead when compaction is not needed. This keeps the tool loop simple (always call `compact()`) while letting each strategy define its own trigger logic.
+
+The following example illustrates this for Variant A (in-place flat list). See Variant C1/C2 under Option 1 for group-aware equivalents.
+
+```python
+class SlidingWindowStrategy(CompactionStrategy):
+ """Example with built-in trigger logic and atomic group preservation (Variant A)."""
+
+ def __init__(self, max_messages: int, *, compact_to: int | None = None):
+ self.max_messages = max_messages
+ self.compact_to = compact_to or max_messages // 2
+
+ async def compact(self, messages: list[ChatMessage]) -> bool:
+ # Fast short-circuit: no-op if under threshold
+ if len(messages) <= self.max_messages:
+ return False
+
+ # Partition into anchors (system messages) and the rest
+ anchors: list[ChatMessage] = []
+ rest: list[ChatMessage] = []
+ for m in messages:
+ (anchors if m.role == "system" else rest).append(m)
+
+ # Group into atomic units: [assistant w/ tool_calls + tool results]
+ # count as one group; standalone messages are their own group
+ groups: list[list[ChatMessage]] = []
+ i = 0
+ while i < len(rest):
+ msg = rest[i]
+ if msg.role == "assistant" and getattr(msg, "tool_calls", None):
+ # Collect this assistant message + all following tool results
+ group = [msg]
+ i += 1
+ while i < len(rest) and rest[i].role == "tool":
+ group.append(rest[i])
+ i += 1
+ groups.append(group)
+ else:
+ groups.append([msg])
+ i += 1
+
+ # Keep the last N groups (by message count) that fit within compact_to
+ kept: list[ChatMessage] = []
+ count = 0
+ for group in reversed(groups):
+ if count + len(group) > self.compact_to:
+ break
+ kept = group + kept
+ count += len(group)
+
+ # Mutate in place
+ messages.clear()
+ messages.extend(anchors + kept)
+ return True
+```
+
+#### Compaction on pre-write and in-run
+
+Given a situation where a compaction strategy is known, the following would need to happen:
+1. At that moment in the run, the message list is passed to the strategy's `compact()` method, which returns whether compaction occurred (and depending on the variant, either mutates in place or returns a new list).
+1. The caller continues with the (potentially reduced) list for the next steps (sending to the model, saving to storage, or continuing the tool loop with the reduced context)
+1. We need to decide how to handle a failed compaction (e.g., the strategy raises an exception) — likely we should have a fallback to continue without compaction rather than failing the entire agent run.
+
+#### Compaction on existing storage
+
+ADR-0016's `HistoryProvider.save_messages()` is an **append** operation — `after_run` collects the new messages from the current invocation and appends them to storage. There is no built-in way to **replace** the full stored history with a compacted version.
+
+For compaction on existing storage (and pre-write compaction that rewrites history), we need a way to overwrite rather than append. Two options:
+
+1. **Add a `replace_messages()` method** to `HistoryProvider`:
+
+```python
+class HistoryProvider(ContextProvider):
+ @abstractmethod
+ async def save_messages(self, session_id: str | None, messages: Sequence[ChatMessage]) -> None:
+ """Append messages to storage for this session."""
+ ...
+
+ async def replace_messages(self, session_id: str | None, messages: Sequence[ChatMessage]) -> None:
+ """Replace all stored messages for this session. Used for compaction.
+
+ Default implementation raises NotImplementedError. Providers that support
+ compaction on existing storage must override this method.
+ """
+ raise NotImplementedError(
+ f"{type(self).__name__} does not support replace_messages. "
+ "Override this method to enable storage compaction."
+ )
+```
+
+2. **Add a `overwrite` parameter** to `save_messages()`:
+
+```python
+class HistoryProvider(ContextProvider):
+ @abstractmethod
+ async def save_messages(
+ self,
+ session_id: str | None,
+ messages: Sequence[ChatMessage],
+ *,
+ overwrite: bool = False,
+ ) -> None:
+ """Persist messages for this session.
+
+ Args:
+ overwrite: If True, replace all existing messages instead of appending.
+ Used for compaction workflows.
+ """
+ ...
+```
+
+Either approach enables the compaction-on-existing-storage workflow:
+
+```python
+history = await provider.get_messages(session_id)
+compacted = await strategy.compact(history)
+await provider.replace_messages(session_id, compacted) # Option 1
+# or
+await provider.save_messages(session_id, compacted, overwrite=True) # Option 2
+```
+
+This could then be combined with a convenience method on the provider for compaction:
+
+```python
+
+class HistoryProvider:
+
+ compaction_strategy: CompactionStrategy | None = None # Optional default strategy for this provider
+
+ async def compact_storage(self, session_id: str | None, *, strategy: CompactionStrategy | None = None) -> None:
+ """Compact stored history for this session using the given strategy."""
+ history = await self.get_messages(session_id)
+ used_strategy = strategy or self._get_strategy("existing") or self._get_strategy("pre_write")
+ if used_strategy is None:
+ raise ValueError("No compaction strategy configured for existing storage.")
+ await used_strategy.compact(history)
+ await self.replace_messages(session_id, history) # or save_messages with overwrite
+ # or
+ await self.save_messages(session_id, history, overwrite=True)
+```
+
+This design choice is orthogonal to the compaction strategy options below — any option requires one of these `HistoryProvider` extensions and optionally the convenience method.
+
+## More Information
+
+### Message Attribution and Compaction
+
+The `source_id` attribution system from ADR-0016 enables intelligent compaction:
+
+```python
+class AttributionAwareStrategy(CompactionStrategy):
+ """Example: remove ephemeral context but preserve RAG and user messages."""
+
+ async def compact(self, messages: list[ChatMessage]) -> bool:
+ ephemeral = [m for m in messages if m.additional_properties.get("source_id") == "ephemeral"]
+ if not ephemeral:
+ return False
+ for msg in ephemeral:
+ messages.remove(msg)
+ return True
+```
+
+### Related Decisions
+
+- [ADR-0016: Unifying Context Management with ContextPlugin](0016-python-context-middleware.md) — Parent ADR that established `ContextProvider`, `HistoryProvider`, and `AgentSession` architecture.
+- [Context Compaction Limitations Analysis](https://gist.github.com/victordibia/ec3f3baf97345f7e47da025cf55b999f) — Detailed analysis of why current architecture cannot support in-run compaction, with attempted solutions and their failure modes. Option 4 in this ADR corresponds to "Option A: Middleware Access to Mutable Message Source" from that analysis; Options 1-3 correspond to "Option B: Tool Loop Hook", adapted here to a `BaseChatClient` hook instead of `FunctionInvocationConfiguration`.
diff --git a/dotnet/Directory.Packages.props b/dotnet/Directory.Packages.props
index 9c308c85fb..255d8fe94f 100644
--- a/dotnet/Directory.Packages.props
+++ b/dotnet/Directory.Packages.props
@@ -19,8 +19,8 @@
-
-
+
+
@@ -35,7 +35,7 @@
-
+
@@ -94,7 +94,7 @@
-
+
@@ -185,4 +185,4 @@
runtime; build; native; contentfiles; analyzers; buildtransitive
-
\ No newline at end of file
+
diff --git a/dotnet/samples/01-get-started/04_memory/Program.cs b/dotnet/samples/01-get-started/04_memory/Program.cs
index 3705e64f3a..a97941620f 100644
--- a/dotnet/samples/01-get-started/04_memory/Program.cs
+++ b/dotnet/samples/01-get-started/04_memory/Program.cs
@@ -89,6 +89,7 @@ namespace SampleApp
internal sealed class UserInfoMemory : AIContextProvider
{
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
private readonly IChatClient _chatClient;
public UserInfoMemory(IChatClient chatClient, Func? stateInitializer = null)
@@ -99,7 +100,7 @@ namespace SampleApp
this._chatClient = chatClient;
}
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
public UserInfo GetUserInfo(AgentSession session)
=> this._sessionState.GetOrInitializeState(session);
diff --git a/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs b/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
index 63fa5c0751..78a8952082 100644
--- a/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
+++ b/dotnet/samples/02-agents/Agents/Agent_Step04_3rdPartyChatHistoryStorage/Program.cs
@@ -79,6 +79,7 @@ namespace SampleApp
internal sealed class VectorChatHistoryProvider : ChatHistoryProvider
{
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
private readonly VectorStore _vectorStore;
public VectorChatHistoryProvider(
@@ -92,7 +93,7 @@ namespace SampleApp
this._vectorStore = vectorStore ?? throw new ArgumentNullException(nameof(vectorStore));
}
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
public string GetSessionDbKey(AgentSession session)
=> this._sessionState.GetOrInitializeState(session).SessionDbKey;
diff --git a/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Evaluations_Step01_RedTeaming/Program.cs b/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Evaluations_Step01_RedTeaming/Program.cs
index 60a859c28f..1e1e48d54b 100644
--- a/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Evaluations_Step01_RedTeaming/Program.cs
+++ b/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Evaluations_Step01_RedTeaming/Program.cs
@@ -60,7 +60,7 @@ Console.WriteLine();
// Submit the red team run to the service
Console.WriteLine("Submitting red team run...");
-RedTeam redTeamRun = await aiProjectClient.RedTeams.CreateAsync(redTeamConfig);
+RedTeam redTeamRun = await aiProjectClient.RedTeams.CreateAsync(redTeamConfig, options: null);
Console.WriteLine($"Red team run created: {redTeamRun.Name}");
Console.WriteLine($"Status: {redTeamRun.Status}");
diff --git a/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Step22_MemorySearch/Program.cs b/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Step22_MemorySearch/Program.cs
index 97eed4e838..836bf1b684 100644
--- a/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Step22_MemorySearch/Program.cs
+++ b/dotnet/samples/02-agents/FoundryAgents/FoundryAgents_Step22_MemorySearch/Program.cs
@@ -35,7 +35,7 @@ string userScope = $"user_{Environment.MachineName}";
AIProjectClient aiProjectClient = new(new Uri(endpoint), new AzureCliCredential());
// Create the Memory Search tool configuration
-MemorySearchTool memorySearchTool = new(memoryStoreName, userScope)
+MemorySearchPreviewTool memorySearchTool = new(memoryStoreName, userScope)
{
// Optional: Configure how quickly new memories are indexed (in seconds)
UpdateDelay = 1,
diff --git a/dotnet/samples/03-workflows/Declarative/HostedWorkflow/Program.cs b/dotnet/samples/03-workflows/Declarative/HostedWorkflow/Program.cs
index 272e83f983..81e2abbafe 100644
--- a/dotnet/samples/03-workflows/Declarative/HostedWorkflow/Program.cs
+++ b/dotnet/samples/03-workflows/Declarative/HostedWorkflow/Program.cs
@@ -88,7 +88,9 @@ internal sealed class Program
{
string workflowYaml = File.ReadAllText("MathChat.yaml");
+#pragma warning disable AAIP001 // WorkflowAgentDefinition is experimental
WorkflowAgentDefinition workflowAgentDefinition = WorkflowAgentDefinition.FromYaml(workflowYaml);
+#pragma warning restore AAIP001
return
await agentClient.CreateAgentAsync(
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
index 82e5f2c360..5ccf139363 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/AIContextProvider.cs
@@ -36,6 +36,8 @@ public abstract class AIContextProvider
private static IEnumerable DefaultNoopFilter(IEnumerable messages)
=> messages;
+ private IReadOnlyList? _stateKeys;
+
///
/// Initializes a new instance of the class.
///
@@ -68,14 +70,15 @@ public abstract class AIContextProvider
protected Func, IEnumerable> StoreInputResponseMessageFilter { get; }
///
- /// Gets the key used to store the provider state in the .
+ /// Gets the set of keys used to store the provider state in the .
///
///
- /// The default value is the name of the concrete type (e.g. "TextSearchProvider").
- /// Implementations may override this to provide a custom key, for example when multiple
- /// instances of the same provider type are used in the same session.
+ /// The default value is a single-element set containing the name of the concrete type (e.g. "TextSearchProvider").
+ /// Implementations may override this to provide custom keys, for example when multiple
+ /// instances of the same provider type are used in the same session, or when a provider
+ /// stores state under more than one key.
///
- public virtual string StateKey => this.GetType().Name;
+ public virtual IReadOnlyList StateKeys => this._stateKeys ??= [this.GetType().Name];
///
/// Called at the start of agent invocation to provide additional context.
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
index df9ff0069e..c7dfb4a233 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/ChatHistoryProvider.cs
@@ -45,6 +45,7 @@ public abstract class ChatHistoryProvider
private static IEnumerable DefaultNoopFilter(IEnumerable messages)
=> messages;
+ private IReadOnlyList? _stateKeys;
private readonly Func, IEnumerable>? _provideOutputMessageFilter;
private readonly Func, IEnumerable> _storeInputRequestMessageFilter;
private readonly Func, IEnumerable> _storeInputResponseMessageFilter;
@@ -66,14 +67,15 @@ public abstract class ChatHistoryProvider
}
///
- /// Gets the key used to store the provider state in the .
+ /// Gets the set of keys used to store the provider state in the .
///
///
- /// The default value is the name of the concrete type (e.g. "InMemoryChatHistoryProvider").
- /// Implementations may override this to provide a custom key, for example when multiple
- /// instances of the same provider type are used in the same session.
+ /// The default value is a single-element set containing the name of the concrete type (e.g. "InMemoryChatHistoryProvider").
+ /// Implementations may override this to provide custom keys, for example when multiple
+ /// instances of the same provider type are used in the same session, or when a provider
+ /// stores state under more than one key.
///
- public virtual string StateKey => this.GetType().Name;
+ public virtual IReadOnlyList StateKeys => this._stateKeys ??= [this.GetType().Name];
///
/// Called at the start of agent invocation to provide messages for the next agent invocation.
diff --git a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
index e09dd6b0a0..7c7b28b7bd 100644
--- a/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Abstractions/InMemoryChatHistoryProvider.cs
@@ -27,6 +27,7 @@ namespace Microsoft.Agents.AI;
public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
{
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
///
/// Initializes a new instance of the class.
@@ -50,7 +51,7 @@ public sealed class InMemoryChatHistoryProvider : ChatHistoryProvider
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
///
/// Gets the chat reducer used to process or reduce chat messages. If null, no reduction logic will be applied.
diff --git a/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs b/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs
index a190f4b154..5d2c67695f 100644
--- a/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs
+++ b/dotnet/src/Microsoft.Agents.AI.AzureAI/AzureAIProjectChatClientExtensions.cs
@@ -39,7 +39,7 @@ public static partial class AzureAIProjectChatClientExtensions
/// The agent with the specified name was not found.
///
/// When instantiating a by using an , minimal information will be available about the agent in the instance level, and any logic that relies
- /// on to retrieve information about the agent like will receive as the result.
+ /// on to retrieve information about the agent like will receive as the result.
///
public static ChatClientAgent AsAIAgent(
this AIProjectClient aiProjectClient,
@@ -355,28 +355,27 @@ public static partial class AzureAIProjectChatClientExtensions
private static readonly ModelReaderWriterOptions s_modelWriterOptionsWire = new("W");
///
- /// Asynchronously retrieves an agent record by name using the Protocol method with user-agent header.
+ /// Asynchronously retrieves an agent record by name using the protocol method to inject user-agent headers.
///
private static async Task GetAgentRecordByNameAsync(AIProjectClient aiProjectClient, string agentName, CancellationToken cancellationToken)
{
ClientResult protocolResponse = await aiProjectClient.Agents.GetAgentAsync(agentName, cancellationToken.ToRequestOptions(false)).ConfigureAwait(false);
var rawResponse = protocolResponse.GetRawResponse();
AgentRecord? result = ModelReaderWriter.Read(rawResponse.Content, s_modelWriterOptionsWire, AzureAIProjectsOpenAIContext.Default);
- return ClientResult.FromOptionalValue(result, rawResponse).Value!
- ?? throw new InvalidOperationException($"Agent with name '{agentName}' not found.");
+ return result ?? throw new InvalidOperationException($"Agent with name '{agentName}' not found.");
}
///
- /// Asynchronously creates an agent version using the Protocol method with user-agent header.
+ /// Asynchronously creates an agent version using the protocol method to inject user-agent headers.
///
private static async Task CreateAgentVersionWithProtocolAsync(AIProjectClient aiProjectClient, string agentName, AgentVersionCreationOptions creationOptions, CancellationToken cancellationToken)
{
- using BinaryContent protocolRequest = BinaryContent.Create(ModelReaderWriter.Write(creationOptions, ModelReaderWriterOptions.Json, AzureAIProjectsContext.Default));
- ClientResult protocolResponse = await aiProjectClient.Agents.CreateAgentVersionAsync(agentName, protocolRequest, cancellationToken.ToRequestOptions(false)).ConfigureAwait(false);
-
+ BinaryData serializedOptions = ModelReaderWriter.Write(creationOptions, s_modelWriterOptionsWire, AzureAIProjectsContext.Default);
+ BinaryContent content = BinaryContent.Create(serializedOptions);
+ ClientResult protocolResponse = await aiProjectClient.Agents.CreateAgentVersionAsync(agentName, content, foundryFeatures: null, cancellationToken.ToRequestOptions(false)).ConfigureAwait(false);
var rawResponse = protocolResponse.GetRawResponse();
AgentVersion? result = ModelReaderWriter.Read(rawResponse.Content, s_modelWriterOptionsWire, AzureAIProjectsOpenAIContext.Default);
- return ClientResult.FromValue(result, rawResponse).Value!;
+ return result ?? throw new InvalidOperationException($"Failed to create agent version for agent '{agentName}'.");
}
private static async Task CreateAIAgentAsync(
diff --git a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
index afaa59ee53..c9238889c9 100644
--- a/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.CosmosNoSql/CosmosChatHistoryProvider.cs
@@ -22,6 +22,7 @@ namespace Microsoft.Agents.AI;
public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
{
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
private readonly CosmosClient _cosmosClient;
private readonly Container _container;
private readonly bool _ownsClient;
@@ -114,7 +115,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
///
/// Initializes a new instance of the class using a connection string.
diff --git a/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
index 0f7041e834..35baa055d1 100644
--- a/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.FoundryMemory/FoundryMemoryProvider.cs
@@ -32,6 +32,7 @@ public sealed class FoundryMemoryProvider : AIContextProvider
private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
private readonly string _contextPrompt;
private readonly string _memoryStoreName;
private readonly int _maxMemories;
@@ -82,7 +83,7 @@ public sealed class FoundryMemoryProvider : AIContextProvider
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
private static Func ValidateStateInitializer(Func stateInitializer) =>
session =>
diff --git a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
index 1e325b5683..678905e395 100644
--- a/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Mem0/Mem0Provider.cs
@@ -27,6 +27,7 @@ public sealed class Mem0Provider : MessageAIContextProvider
private const string DefaultContextPrompt = "## Memories\nConsider the following memories when answering user questions:";
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
private readonly string _contextPrompt;
private readonly bool _enableSensitiveTelemetryData;
@@ -72,7 +73,7 @@ public sealed class Mem0Provider : MessageAIContextProvider
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
private static Func ValidateStateInitializer(Func stateInitializer) =>
session =>
diff --git a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
index 1fd42f923e..2815ed99f0 100644
--- a/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI.Workflows/WorkflowChatHistoryProvider.cs
@@ -12,6 +12,7 @@ namespace Microsoft.Agents.AI.Workflows;
internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
{
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
///
/// Initializes a new instance of the class.
@@ -30,7 +31,7 @@ internal sealed class WorkflowChatHistoryProvider : ChatHistoryProvider
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
internal sealed class StoreState
{
diff --git a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs
index d52ea52e43..7db4eff6d8 100644
--- a/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs
+++ b/dotnet/src/Microsoft.Agents.AI/ChatClient/ChatClientAgent.cs
@@ -112,7 +112,7 @@ public sealed partial class ChatClientAgent : AIAgent
this.ChatHistoryProvider = options?.ChatHistoryProvider ?? new InMemoryChatHistoryProvider();
this.AIContextProviders = this._agentOptions?.AIContextProviders as IReadOnlyList ?? this._agentOptions?.AIContextProviders?.ToList();
- // Validate that no two providers share the same StateKey, since they would overwrite each other's state in the session.
+ // Validate that no two providers share any StateKeys, since they would overwrite each other's state in the session.
this._aiContextProviderStateKeys = ValidateAndCollectStateKeys(this._agentOptions?.AIContextProviders, this.ChatHistoryProvider);
this._logger = (loggerFactory ?? chatClient.GetService() ?? NullLoggerFactory.Instance).CreateLogger();
@@ -824,11 +824,17 @@ public sealed partial class ChatClientAgent : AIAgent
$"Only {nameof(ChatClientAgentSession.ConversationId)} or {nameof(this.ChatHistoryProvider)} may be used, but not both. The current {nameof(ChatClientAgentSession)} has a {nameof(ChatClientAgentSession.ConversationId)} indicating server-side chat history management, but an override {nameof(this.ChatHistoryProvider)} was provided via {nameof(AgentRunOptions.AdditionalProperties)}.");
}
- // Validate that the override provider's StateKey does not clash with any AIContextProvider's StateKey.
- if (overrideProvider is not null && this._aiContextProviderStateKeys.Contains(overrideProvider.StateKey))
+ // Validate that the override provider's StateKeys do not clash with any AIContextProvider's StateKeys.
+ if (overrideProvider is not null)
{
- throw new InvalidOperationException(
- $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses the state key '{overrideProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state.");
+ foreach (var key in overrideProvider.StateKeys)
+ {
+ if (this._aiContextProviderStateKeys.Contains(key))
+ {
+ throw new InvalidOperationException(
+ $"The ChatHistoryProvider '{overrideProvider.GetType().Name}' uses state key '{key}' which is already used by one of the configured AIContextProviders. Each provider must use unique state keys to avoid overwriting each other's state.");
+ }
+ }
}
provider = overrideProvider;
@@ -879,7 +885,7 @@ public sealed partial class ChatClientAgent : AIAgent
private string GetLoggingAgentName() => this.Name ?? "UnnamedAgent";
///
- /// Validates that all configured providers have unique values
+ /// Validates that all configured providers have unique values
/// and returns a of the AIContextProvider state keys.
///
private static HashSet ValidateAndCollectStateKeys(IEnumerable? aiContextProviders, ChatHistoryProvider? chatHistoryProvider)
@@ -890,10 +896,13 @@ public sealed partial class ChatClientAgent : AIAgent
{
foreach (var provider in aiContextProviders)
{
- if (!stateKeys.Add(provider.StateKey))
+ foreach (var key in provider.StateKeys)
{
- throw new InvalidOperationException(
- $"Multiple providers use the same state key '{provider.StateKey}'. Each provider must use a unique state key to avoid overwriting each other's state.");
+ if (!stateKeys.Add(key))
+ {
+ throw new InvalidOperationException(
+ $"Multiple providers use the same state key '{key}'. Each provider must use a unique state key to avoid overwriting each other's state.");
+ }
}
}
}
@@ -905,11 +914,16 @@ public sealed partial class ChatClientAgent : AIAgent
$"The default {nameof(InMemoryChatHistoryProvider)} uses the state key '{nameof(InMemoryChatHistoryProvider)}', which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{nameof(InMemoryChatHistoryProvider)}' as its state key, or provide a custom ChatHistoryProvider with a unique state key.");
}
- if (chatHistoryProvider is not null
- && stateKeys.Contains(chatHistoryProvider.StateKey))
+ if (chatHistoryProvider is not null)
{
- throw new InvalidOperationException(
- $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses the state key '{chatHistoryProvider.StateKey}' which is already used by one of the configured AIContextProviders. Each provider must use a unique state key to avoid overwriting each other's state. To resolve this, either configure a different state key for the AIContextProvider that is using '{chatHistoryProvider.StateKey}' as its state key, or reconfigure the custom ChatHistoryProvider with a unique state key.");
+ foreach (var key in chatHistoryProvider.StateKeys)
+ {
+ if (stateKeys.Contains(key))
+ {
+ throw new InvalidOperationException(
+ $"The ChatHistoryProvider '{chatHistoryProvider.GetType().Name}' uses state key '{key}' which is already used by one of the configured AIContextProviders. Each provider must use unique state keys to avoid overwriting each other's state. To resolve this, either configure different state keys for the AIContextProvider that shares keys with the ChatHistoryProvider, or reconfigure the custom ChatHistoryProvider with unique state keys.");
+ }
+ }
}
return stateKeys;
diff --git a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
index cd59d1aaa3..80d5e1144f 100644
--- a/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI/Memory/ChatHistoryMemoryProvider.cs
@@ -54,6 +54,7 @@ public sealed class ChatHistoryMemoryProvider : MessageAIContextProvider, IDispo
private const string ContentEmbeddingField = "ContentEmbedding";
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
#pragma warning disable CA2213 // VectorStore is not owned by this class - caller is responsible for disposal
private readonly VectorStore _vectorStore;
@@ -128,7 +129,7 @@ public sealed class ChatHistoryMemoryProvider : MessageAIContextProvider, IDispo
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
///
protected override async ValueTask ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)
diff --git a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
index df53729fce..11611f0f69 100644
--- a/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
+++ b/dotnet/src/Microsoft.Agents.AI/TextSearchProvider.cs
@@ -40,6 +40,7 @@ public sealed class TextSearchProvider : MessageAIContextProvider
private const string DefaultCitationsPrompt = "Include citations to the source document with document name and link if document name and link is available.";
private readonly ProviderSessionState _sessionState;
+ private IReadOnlyList? _stateKeys;
private readonly Func>> _searchAsync;
private readonly ILogger? _logger;
private readonly AITool[] _tools;
@@ -88,7 +89,7 @@ public sealed class TextSearchProvider : MessageAIContextProvider
}
///
- public override string StateKey => this._sessionState.StateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys ??= [this._sessionState.StateKey];
///
protected override async ValueTask ProvideAIContextAsync(AIContextProvider.InvokingContext context, CancellationToken cancellationToken = default)
diff --git a/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentCreateTests.cs b/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentCreateTests.cs
index ab2e1848a5..f750b5a8e7 100644
--- a/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentCreateTests.cs
+++ b/dotnet/tests/AzureAIAgentsPersistent.IntegrationTests/AzureAIAgentsPersistentCreateTests.cs
@@ -132,10 +132,15 @@ public class AzureAIAgentsPersistentCreateTests
}
}
- [Theory]
- [InlineData("CreateWithChatClientAgentOptionsAsync")]
- [InlineData("CreateWithFoundryOptionsAsync")]
- public async Task CreateAgent_CreatesAgentWithCodeInterpreterAsync(string createMechanism)
+ [Fact]
+ public Task CreateAgent_CreatesAgentWithCodeInterpreter_ChatClientAgentOptionsAsync()
+ => this.CreateAgent_CreatesAgentWithCodeInterpreterAsync("CreateWithChatClientAgentOptionsAsync");
+
+ [RetryFact(Constants.RetryCount, Constants.RetryDelay)]
+ public Task CreateAgent_CreatesAgentWithCodeInterpreter_FoundryOptionsAsync()
+ => this.CreateAgent_CreatesAgentWithCodeInterpreterAsync("CreateWithFoundryOptionsAsync");
+
+ private async Task CreateAgent_CreatesAgentWithCodeInterpreterAsync(string createMechanism)
{
// Arrange.
const string AgentInstructions = """
diff --git a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs
index 147ceaf195..94beb08bdf 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Abstractions.UnitTests/InMemoryChatHistoryProviderTests.cs
@@ -43,23 +43,25 @@ public class InMemoryChatHistoryProviderTests
}
[Fact]
- public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+ public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
{
// Arrange & Act
var provider = new InMemoryChatHistoryProvider();
// Assert
- Assert.Equal("InMemoryChatHistoryProvider", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("InMemoryChatHistoryProvider", provider.StateKeys);
}
[Fact]
- public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+ public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
{
// Arrange & Act
var provider = new InMemoryChatHistoryProvider(new() { StateKey = "custom-key" });
// Assert
- Assert.Equal("custom-key", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("custom-key", provider.StateKeys);
}
[Fact]
diff --git a/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs b/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs
index a7b9c54aac..65726bb2aa 100644
--- a/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientExtensionsTests.cs
@@ -467,7 +467,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithModelAndOptions_CreatesValidAgentAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", instructions: "Test instructions");
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", instructions: "Test instructions");
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -475,7 +475,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- var agent = await client.CreateAIAgentAsync("test-model", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -490,7 +490,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithModelAndOptions_WithClientFactory_AppliesFactoryCorrectlyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", instructions: "Test instructions");
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", instructions: "Test instructions");
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -499,7 +499,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
TestChatClient? testChatClient = null;
// Act
- var agent = await client.CreateAIAgentAsync(
+ var agent = await testClient.Client.CreateAIAgentAsync(
"test-model",
options,
clientFactory: (innerClient) => testChatClient = new TestChatClient(innerClient));
@@ -560,12 +560,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithDefinition_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -582,12 +582,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
var definitionResponse = GeneratePromptDefinitionResponse(definition, null);
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -602,12 +602,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
{
// Arrange
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definition);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definition);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -628,12 +628,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
// Create a response definition with the same tool
var definitionResponse = GeneratePromptDefinitionResponse(definition, definition.Tools.Select(t => t.AsAITool()).ToList());
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -667,12 +667,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
definitionResponse.Tools.Add(tool);
}
- AIProjectClient client = this.CreateTestAgentClient(agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -803,10 +803,10 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var definitionResponse = GeneratePromptDefinitionResponse(new PromptAgentDefinition("test-model") { Instructions = "Test instructions" }, tools);
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
// Act
- var agent = await client.CreateAIAgentAsync(
+ var agent = await testClient.Client.CreateAIAgentAsync(
"test-agent",
"test-model",
"Test instructions",
@@ -831,14 +831,14 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithDefinitionTools_CreatesAgentAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test instructions" };
definition.Tools.Add(ResponseTool.CreateFunctionTool("async_tool", BinaryData.FromString("{}"), strictModeEnabled: false));
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -885,7 +885,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var sharepointOptions = new SharePointGroundingToolOptions();
sharepointOptions.ProjectConnections.Add(new ToolProjectConnection("connection-id"));
- var structuredOutputs = new StructuredOutputDefinition("name", "description", BinaryData.FromString(AIJsonUtilities.CreateJsonSchema(new { id = "test" }.GetType()).ToString()), false);
+ var structuredOutputs = new StructuredOutputDefinition("name", "description", new Dictionary { ["schema"] = BinaryData.FromString(AIJsonUtilities.CreateJsonSchema(new { id = "test" }.GetType()).ToString()) }, false);
// Add tools to the definition
definition.Tools.Add(ResponseTool.CreateFunctionTool("create_tool", BinaryData.FromString("{}"), strictModeEnabled: false));
@@ -902,12 +902,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
// Generate agent definition response with the tools
var definitionResponse = GeneratePromptDefinitionResponse(definition, definition.Tools.Select(t => t.AsAITool()).ToList());
- AIProjectClient client = this.CreateTestAgentClient(agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -942,12 +942,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var definitionResponse = new PromptAgentDefinition("test-model") { Instructions = "Test" };
definitionResponse.Tools.Add(functionTool);
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -961,7 +961,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithDeclarativeFunctionFromDefinition_AcceptsDeclarativeFunctionAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
// Create a declarative function (not invocable) using AIFunctionFactory.CreateDeclaration
@@ -974,7 +974,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -1001,12 +1001,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var definitionResponse = new PromptAgentDefinition("test-model") { Instructions = "Test" };
definitionResponse.Tools.Add(declarativeFunction.AsOpenAIResponseTool() ?? throw new InvalidOperationException());
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -1027,12 +1027,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test instructions" };
var definitionResponse = GeneratePromptDefinitionResponse(definition, null);
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync("test-agent", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-agent", options);
// Assert
Assert.NotNull(agent);
@@ -1083,7 +1083,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
new PromptAgentDefinition("test-model") { Instructions = "Test" },
tools);
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: definitionResponse);
var options = new ChatClientAgentOptions
{
@@ -1092,7 +1092,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- var agent = await client.CreateAIAgentAsync("test-model", options);
+ var agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -1278,14 +1278,14 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithClientFactory_ReceivesCorrectUnderlyingClientAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
IChatClient? receivedClient = null;
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync(
+ var agent = await testClient.Client.CreateAIAgentAsync(
"test-agent",
options,
clientFactory: (innerClient) =>
@@ -1340,10 +1340,10 @@ public sealed class AzureAIProjectChatClientExtensionsTests
const string AgentName = "test-agent";
const string Model = "test-model";
const string Instructions = "Test instructions";
- AIProjectClient client = this.CreateTestAgentClient(AgentName, Instructions);
+ using var testClient = CreateTestAgentClientWithHandler(AgentName, Instructions);
// Act
- var agent = await client.CreateAIAgentAsync(
+ var agent = await testClient.Client.CreateAIAgentAsync(
AgentName,
Model,
Instructions,
@@ -1367,12 +1367,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
var definition = new PromptAgentDefinition("test-model") { Instructions = "Test" };
var agentDefinitionResponse = GeneratePromptDefinitionResponse(definition, null);
- AIProjectClient client = this.CreateTestAgentClient(agentName: "test-agent", agentDefinitionResponse: agentDefinitionResponse);
+ using var testClient = CreateTestAgentClientWithHandler(agentName: "test-agent", agentDefinitionResponse: agentDefinitionResponse);
var options = new AgentVersionCreationOptions(definition);
// Act
- var agent = await client.CreateAIAgentAsync(
+ var agent = await testClient.Client.CreateAIAgentAsync(
"test-agent",
options,
clientFactory: (innerClient) => new TestChatClient(innerClient));
@@ -1390,7 +1390,8 @@ public sealed class AzureAIProjectChatClientExtensionsTests
#region User-Agent Header Tests
///
- /// Verifies that the user-agent header is added to both synchronous and asynchronous requests made by agent creation methods.
+ /// Verifies that the MEAI user-agent header is added to CreateAIAgentAsync POST requests
+ /// via the protocol method's RequestOptions pipeline policy.
///
[Fact]
public async Task CreateAIAgentAsync_UserAgentHeaderAddedToRequestsAsync()
@@ -1398,9 +1399,12 @@ public sealed class AzureAIProjectChatClientExtensionsTests
using var httpHandler = new HttpHandlerAssert(request =>
{
Assert.Equal("POST", request.Method.Method);
- Assert.Contains("MEAI", request.Headers.UserAgent.ToString());
- return new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(TestDataUtil.GetAgentResponseJson(), Encoding.UTF8, "application/json") };
+ // Verify MEAI user-agent header is present on CreateAgentVersion POST request
+ Assert.True(request.Headers.TryGetValues("User-Agent", out var userAgentValues));
+ Assert.Contains(userAgentValues, v => v.Contains("MEAI"));
+
+ return new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(TestDataUtil.GetAgentVersionResponseJson(), Encoding.UTF8, "application/json") };
});
#pragma warning disable CA5399
@@ -1940,7 +1944,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithTextResponseFormat_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -1952,7 +1956,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -1966,7 +1970,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithJsonResponseFormatWithoutSchema_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -1978,7 +1982,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -1992,7 +1996,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithJsonResponseFormatWithSchema_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
JsonElement schemaElement = AIJsonUtilities.CreateJsonSchema(typeof(TestSchema));
var jsonFormat = ChatResponseFormat.ForJsonSchema(schemaElement, "test_schema", "A test schema");
var options = new ChatClientAgentOptions
@@ -2006,7 +2010,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2020,7 +2024,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithJsonResponseFormatWithSchemaAndStrictMode_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
JsonElement schemaElement = AIJsonUtilities.CreateJsonSchema(typeof(TestSchema));
var jsonFormat = ChatResponseFormat.ForJsonSchema(schemaElement, "test_schema", "A test schema");
var additionalProps = new AdditionalPropertiesDictionary
@@ -2039,7 +2043,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2053,7 +2057,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithJsonResponseFormatWithSchemaAndStrictModeFalse_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
JsonElement schemaElement = AIJsonUtilities.CreateJsonSchema(typeof(TestSchema));
var jsonFormat = ChatResponseFormat.ForJsonSchema(schemaElement, "test_schema", "A test schema");
var additionalProps = new AdditionalPropertiesDictionary
@@ -2072,7 +2076,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2090,7 +2094,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithRawRepresentationFactory_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -2102,7 +2106,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2116,7 +2120,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithRawRepresentationFactoryReturningNull_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -2128,7 +2132,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2142,7 +2146,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithRawRepresentationFactoryReturningNonCreateResponseOptions_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -2154,7 +2158,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2172,7 +2176,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithDescription_SetsDescriptionAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient(description: "Test description");
+ using var testClient = CreateTestAgentClientWithHandler(description: "Test description");
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -2181,7 +2185,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2195,7 +2199,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithoutDescription_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var options = new ChatClientAgentOptions
{
Name = "test-agent",
@@ -2203,7 +2207,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2688,7 +2692,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
public async Task CreateAIAgentAsync_WithHostedToolTypes_CreatesAgentSuccessfullyAsync()
{
// Arrange
- AIProjectClient client = this.CreateTestAgentClient();
+ using var testClient = CreateTestAgentClientWithHandler();
var webSearchTool = new HostedWebSearchTool();
var options = new ChatClientAgentOptions
@@ -2702,7 +2706,7 @@ public sealed class AzureAIProjectChatClientExtensionsTests
};
// Act
- ChatClientAgent agent = await client.CreateAIAgentAsync("test-model", options);
+ ChatClientAgent agent = await testClient.Client.CreateAIAgentAsync("test-model", options);
// Assert
Assert.NotNull(agent);
@@ -2855,6 +2859,54 @@ public sealed class AzureAIProjectChatClientExtensionsTests
return new FakeAgentClient(agentName, instructions, description, agentDefinitionResponse);
}
+ ///
+ /// Creates a test AIProjectClient backed by an HTTP handler that returns canned responses.
+ /// Used for tests that exercise the protocol-method code path (CreateAgentVersion).
+ /// The returned client must be disposed to clean up the underlying HttpClient/handler.
+ ///
+ private static DisposableTestClient CreateTestAgentClientWithHandler(string? agentName = null, string? instructions = null, string? description = null, AgentDefinition? agentDefinitionResponse = null)
+ {
+ var responseJson = TestDataUtil.GetAgentVersionResponseJson(agentName, agentDefinitionResponse, instructions, description);
+
+ var httpHandler = new HttpHandlerAssert(_ =>
+ new HttpResponseMessage(HttpStatusCode.OK) { Content = new StringContent(responseJson, Encoding.UTF8, "application/json") });
+
+#pragma warning disable CA5399
+ var httpClient = new HttpClient(httpHandler);
+#pragma warning restore CA5399
+
+ var client = new AIProjectClient(
+ new Uri("https://test.openai.azure.com/"),
+ new FakeAuthenticationTokenProvider(),
+ new() { Transport = new HttpClientPipelineTransport(httpClient) });
+
+ return new DisposableTestClient(client, httpClient, httpHandler);
+ }
+
+ ///
+ /// Wraps an AIProjectClient and its disposable dependencies for deterministic cleanup.
+ ///
+ private sealed class DisposableTestClient : IDisposable
+ {
+ private readonly HttpClient _httpClient;
+ private readonly HttpHandlerAssert _httpHandler;
+
+ public DisposableTestClient(AIProjectClient client, HttpClient httpClient, HttpHandlerAssert httpHandler)
+ {
+ this.Client = client;
+ this._httpClient = httpClient;
+ this._httpHandler = httpHandler;
+ }
+
+ public AIProjectClient Client { get; }
+
+ public void Dispose()
+ {
+ this._httpClient.Dispose();
+ this._httpHandler.Dispose();
+ }
+ }
+
///
/// Creates a test AgentRecord for testing.
///
@@ -3039,25 +3091,13 @@ public sealed class AzureAIProjectChatClientExtensionsTests
return Task.FromResult(ClientResult.FromValue(ModelReaderWriter.Read(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200)));
}
- public override ClientResult CreateAgentVersion(string agentName, BinaryContent content, RequestOptions? options = null)
- {
- var responseJson = this.GetAgentVersionResponseJson();
- return ClientResult.FromValue(ModelReaderWriter.Read(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200, BinaryData.FromString(responseJson)));
- }
-
- public override ClientResult CreateAgentVersion(string agentName, AgentVersionCreationOptions? options = null, CancellationToken cancellationToken = default)
+ public override ClientResult CreateAgentVersion(string agentName, AgentVersionCreationOptions? options = null, string? foundryFeatures = null, CancellationToken cancellationToken = default)
{
var responseJson = this.GetAgentVersionResponseJson();
return ClientResult.FromValue(ModelReaderWriter.Read(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200));
}
- public override Task CreateAgentVersionAsync(string agentName, BinaryContent content, RequestOptions? options = null)
- {
- var responseJson = this.GetAgentVersionResponseJson();
- return Task.FromResult(ClientResult.FromValue(ModelReaderWriter.Read(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200, BinaryData.FromString(responseJson))));
- }
-
- public override Task> CreateAgentVersionAsync(string agentName, AgentVersionCreationOptions? options = null, CancellationToken cancellationToken = default)
+ public override Task> CreateAgentVersionAsync(string agentName, AgentVersionCreationOptions? options = null, string? foundryFeatures = null, CancellationToken cancellationToken = default)
{
var responseJson = this.GetAgentVersionResponseJson();
return Task.FromResult(ClientResult.FromValue(ModelReaderWriter.Read(BinaryData.FromString(responseJson))!, new MockPipelineResponse(200)));
diff --git a/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientTests.cs
index 9cc340ef5e..5c61e0b457 100644
--- a/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.AzureAI.UnitTests/AzureAIProjectChatClientTests.cs
@@ -22,7 +22,7 @@ public class AzureAIProjectChatClientTests
var requestTriggered = false;
using var httpHandler = new HttpHandlerAssert(async (request) =>
{
- if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+ if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
{
requestTriggered = true;
@@ -71,7 +71,7 @@ public class AzureAIProjectChatClientTests
var requestTriggered = false;
using var httpHandler = new HttpHandlerAssert(async (request) =>
{
- if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+ if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
{
requestTriggered = true;
@@ -120,7 +120,7 @@ public class AzureAIProjectChatClientTests
var requestTriggered = false;
using var httpHandler = new HttpHandlerAssert(async (request) =>
{
- if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+ if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
{
requestTriggered = true;
@@ -169,7 +169,7 @@ public class AzureAIProjectChatClientTests
var requestTriggered = false;
using var httpHandler = new HttpHandlerAssert(async (request) =>
{
- if (request.RequestUri!.PathAndQuery.Contains("openai/responses"))
+ if (request.Method == HttpMethod.Post && request.RequestUri!.PathAndQuery.Contains("/responses"))
{
requestTriggered = true;
diff --git a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs
index 079c096d84..4b62e549c0 100644
--- a/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.CosmosNoSql.UnitTests/CosmosChatHistoryProviderTests.cs
@@ -152,7 +152,7 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
[Fact]
[Trait("Category", "CosmosDB")]
- public void StateKey_ReturnsDefaultKey_WhenNoStateKeyProvided()
+ public void StateKeys_ReturnsDefaultKey_WhenNoStateKeyProvided()
{
// Arrange & Act
this.SkipIfEmulatorNotAvailable();
@@ -161,12 +161,13 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
_ => new CosmosChatHistoryProvider.State("test-conversation"));
// Assert
- Assert.Equal("CosmosChatHistoryProvider", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("CosmosChatHistoryProvider", provider.StateKeys);
}
[Fact]
[Trait("Category", "CosmosDB")]
- public void StateKey_ReturnsCustomKey_WhenSetViaConstructor()
+ public void StateKeys_ReturnsCustomKey_WhenSetViaConstructor()
{
// Arrange & Act
this.SkipIfEmulatorNotAvailable();
@@ -176,7 +177,8 @@ public sealed class CosmosChatHistoryProviderTests : IAsyncLifetime, IDisposable
stateKey: "custom-key");
// Assert
- Assert.Equal("custom-key", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("custom-key", provider.StateKeys);
}
[Fact]
diff --git a/dotnet/tests/Microsoft.Agents.AI.GitHub.Copilot.UnitTests/GitHubCopilotAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.GitHub.Copilot.UnitTests/GitHubCopilotAgentTests.cs
index 8a4d3c1068..5806636925 100644
--- a/dotnet/tests/Microsoft.Agents.AI.GitHub.Copilot.UnitTests/GitHubCopilotAgentTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.GitHub.Copilot.UnitTests/GitHubCopilotAgentTests.cs
@@ -109,7 +109,7 @@ public sealed class GitHubCopilotAgentTests
var hooks = new SessionHooks();
var infiniteSessions = new InfiniteSessionConfig();
var systemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = "Be helpful" };
- PermissionHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
+ PermissionRequestHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
UserInputHandler userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
var mcpServers = new Dictionary { ["server1"] = new McpLocalServerConfig() };
@@ -160,7 +160,7 @@ public sealed class GitHubCopilotAgentTests
var hooks = new SessionHooks();
var infiniteSessions = new InfiniteSessionConfig();
var systemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = "Be helpful" };
- PermissionHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
+ PermissionRequestHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
UserInputHandler userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
var mcpServers = new Dictionary { ["server1"] = new McpLocalServerConfig() };
diff --git a/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs
index 9f9de9127b..3374270861 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Mem0.UnitTests/Mem0ProviderTests.cs
@@ -67,17 +67,18 @@ public sealed class Mem0ProviderTests : IDisposable
}
[Fact]
- public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+ public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
{
// Arrange & Act
var provider = new Mem0Provider(this._httpClient, _ => new Mem0Provider.State(new Mem0ProviderScope { ThreadId = "tid" }));
// Assert
- Assert.Equal("Mem0Provider", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("Mem0Provider", provider.StateKeys);
}
[Fact]
- public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+ public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
{
// Arrange & Act
var provider = new Mem0Provider(
@@ -86,7 +87,8 @@ public sealed class Mem0ProviderTests : IDisposable
new Mem0ProviderOptions { StateKey = "custom-key" });
// Assert
- Assert.Equal("custom-key", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("custom-key", provider.StateKeys);
}
[Fact]
@@ -419,7 +421,7 @@ public sealed class Mem0ProviderTests : IDisposable
}
[Fact]
- public async Task StateKey_CanBeConfiguredViaOptionsAsync()
+ public async Task StateKeys_CanBeConfiguredViaOptionsAsync()
{
// Arrange
this._handler.EnqueueJsonResponse("[]");
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
index 51eb4be3ab..3b06bbb772 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AIContextProviderDecorators/AIContextProviderChatClientTests.cs
@@ -380,7 +380,7 @@ public class AIContextProviderChatClientTests
///
private sealed class TestAIContextProvider : AIContextProvider
{
- private readonly string _stateKey;
+ private readonly IReadOnlyList _stateKeys;
private readonly IEnumerable _provideMessages;
private readonly string? _provideInstructions;
private readonly IEnumerable? _provideTools;
@@ -389,7 +389,7 @@ public class AIContextProviderChatClientTests
public InvokedContext? LastInvokedContext { get; private set; }
- public override string StateKey => this._stateKey;
+ public override IReadOnlyList StateKeys => this._stateKeys;
public TestAIContextProvider(
string stateKey,
@@ -397,7 +397,7 @@ public class AIContextProviderChatClientTests
string? provideInstructions = null,
IEnumerable? provideTools = null)
{
- this._stateKey = stateKey;
+ this._stateKeys = [stateKey];
this._provideMessages = provideMessages ?? [];
this._provideInstructions = provideInstructions;
this._provideTools = provideTools;
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
index 9713a91c2c..2b3cfe43e8 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgentTests.cs
@@ -105,8 +105,8 @@ public partial class ChatClientAgentTests
ChatHistoryProvider = historyProvider
}));
- Assert.Contains("SharedKey", ex.Message);
- Assert.Contains(nameof(ChatHistoryProvider), ex.Message);
+ Assert.Contains("ChatHistoryProvider", ex.Message);
+ Assert.Contains("state key 'SharedKey'", ex.Message);
}
///
@@ -159,11 +159,11 @@ public partial class ChatClientAgentTests
var ex = await Assert.ThrowsAsync(() =>
agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }));
- Assert.Contains("SharedKey", ex.Message);
+ Assert.Contains("state key 'SharedKey'", ex.Message);
}
///
- /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKey as the default ChatHistoryProvider.
+ /// Verify that RunAsync succeeds when an override ChatHistoryProvider uses the same StateKeys as the default ChatHistoryProvider.
///
[Fact]
public async Task RunAsync_SucceedsWhenOverrideChatHistoryProviderSharesKeyWithDefaultAsync()
@@ -192,6 +192,102 @@ public partial class ChatClientAgentTests
await agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties });
}
+ ///
+ /// Verify that the constructor throws when two multi-key AIContextProviders have an overlapping key.
+ ///
+ [Fact]
+ public void Constructor_ThrowsWhenMultiKeyAIContextProvidersOverlap()
+ {
+ // Arrange
+ var chatClient = new Mock().Object;
+ var provider1 = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+ var provider2 = new MultiKeyTestAIContextProvider("Key2", "SharedKey");
+
+ // Act & Assert
+ var ex = Assert.Throws(() =>
+ new ChatClientAgent(chatClient, options: new()
+ {
+ AIContextProviders = [provider1, provider2]
+ }));
+
+ Assert.Contains("state key 'SharedKey'", ex.Message);
+ }
+
+ ///
+ /// Verify that the constructor throws when a multi-key ChatHistoryProvider has an overlapping key with an AIContextProvider.
+ ///
+ [Fact]
+ public void Constructor_ThrowsWhenMultiKeyChatHistoryProviderOverlapsWithAIContextProvider()
+ {
+ // Arrange
+ var chatClient = new Mock().Object;
+ var contextProvider = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+ var historyProvider = new MultiKeyTestChatHistoryProvider("Key2", "SharedKey");
+
+ // Act & Assert
+ var ex = Assert.Throws(() =>
+ new ChatClientAgent(chatClient, options: new()
+ {
+ AIContextProviders = [contextProvider],
+ ChatHistoryProvider = historyProvider
+ }));
+
+ Assert.Contains("state key 'SharedKey'", ex.Message);
+ }
+
+ ///
+ /// Verify that the constructor succeeds when multi-key providers have no overlapping keys.
+ ///
+ [Fact]
+ public void Constructor_SucceedsWithMultiKeyProvidersWithUniqueKeys()
+ {
+ // Arrange
+ var chatClient = new Mock().Object;
+ var contextProvider1 = new MultiKeyTestAIContextProvider("Key1", "Key2");
+ var contextProvider2 = new MultiKeyTestAIContextProvider("Key3", "Key4");
+ var historyProvider = new MultiKeyTestChatHistoryProvider("Key5", "Key6");
+
+ // Act & Assert - should not throw
+ _ = new ChatClientAgent(chatClient, options: new()
+ {
+ AIContextProviders = [contextProvider1, contextProvider2],
+ ChatHistoryProvider = historyProvider
+ });
+ }
+
+ ///
+ /// Verify that RunAsync throws when a multi-key override ChatHistoryProvider has an overlapping key with an AIContextProvider.
+ ///
+ [Fact]
+ public async Task RunAsync_ThrowsWhenMultiKeyOverrideChatHistoryProviderClashesWithAIContextProviderAsync()
+ {
+ // Arrange
+ Mock mockService = new();
+ mockService.Setup(
+ s => s.GetResponseAsync(
+ It.IsAny>(),
+ It.IsAny(),
+ It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
+
+ var contextProvider = new MultiKeyTestAIContextProvider("Key1", "SharedKey");
+ var overrideHistoryProvider = new MultiKeyTestChatHistoryProvider("Key2", "SharedKey");
+
+ ChatClientAgent agent = new(mockService.Object, options: new()
+ {
+ AIContextProviders = [contextProvider]
+ });
+
+ // Act & Assert
+ ChatClientAgentSession? session = await agent.CreateSessionAsync() as ChatClientAgentSession;
+ AdditionalPropertiesDictionary additionalProperties = new();
+ additionalProperties.Add(overrideHistoryProvider);
+
+ var ex = await Assert.ThrowsAsync(() =>
+ agent.RunAsync([new(ChatRole.User, "test")], session, options: new AgentRunOptions { AdditionalProperties = additionalProperties }));
+
+ Assert.Contains("state key 'SharedKey'", ex.Message);
+ }
+
#endregion
#region RunAsync Tests
@@ -489,6 +585,7 @@ public partial class ChatClientAgentTests
.ReturnsAsync(new ChatResponse(responseMessages));
var mockProvider = new Mock(null, null, null);
+ mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
mockProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -560,6 +657,7 @@ public partial class ChatClientAgentTests
.Throws(new InvalidOperationException("downstream failure"));
var mockProvider = new Mock(null, null, null);
+ mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
mockProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -618,6 +716,7 @@ public partial class ChatClientAgentTests
.ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
var mockProvider = new Mock(null, null, null);
+ mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
mockProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -678,7 +777,7 @@ public partial class ChatClientAgentTests
// Provider 1: adds a system message and a tool
var mockProvider1 = new Mock(null, null, null);
- mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockProvider1
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -697,7 +796,7 @@ public partial class ChatClientAgentTests
// Provider 2: adds another system message and verifies it receives accumulated context from provider 1
AIContext? provider2ReceivedContext = null;
var mockProvider2 = new Mock(null, null, null);
- mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+ mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
mockProvider2
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -785,7 +884,7 @@ public partial class ChatClientAgentTests
.ThrowsAsync(new InvalidOperationException("downstream failure"));
var mockProvider1 = new Mock(null, null, null);
- mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockProvider1
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -802,7 +901,7 @@ public partial class ChatClientAgentTests
.Returns(new ValueTask());
var mockProvider2 = new Mock(null, null, null);
- mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+ mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
mockProvider2
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -870,7 +969,7 @@ public partial class ChatClientAgentTests
.Returns(ToAsyncEnumerableAsync(responseUpdates));
var mockProvider1 = new Mock(null, null, null);
- mockProvider1.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockProvider1.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockProvider1
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -887,7 +986,7 @@ public partial class ChatClientAgentTests
.Returns(new ValueTask());
var mockProvider2 = new Mock(null, null, null);
- mockProvider2.SetupGet(p => p.StateKey).Returns("Provider2");
+ mockProvider2.SetupGet(p => p.StateKeys).Returns(["Provider2"]);
mockProvider2
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -1829,6 +1928,7 @@ public partial class ChatClientAgentTests
.Returns(ToAsyncEnumerableAsync(responseUpdates));
var mockProvider = new Mock(null, null, null);
+ mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
mockProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -1908,6 +2008,7 @@ public partial class ChatClientAgentTests
.Throws(new InvalidOperationException("downstream failure"));
var mockProvider = new Mock(null, null, null);
+ mockProvider.SetupGet(p => p.StateKeys).Returns(["TestProvider"]);
mockProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -1965,7 +2066,17 @@ public partial class ChatClientAgentTests
private sealed class TestAIContextProvider(string stateKey) : AIContextProvider
{
- public override string StateKey => stateKey;
+ private readonly IReadOnlyList _stateKeys = [stateKey];
+
+ public override IReadOnlyList StateKeys => this._stateKeys;
+
+ protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+ => new(context.AIContext);
+ }
+
+ private sealed class MultiKeyTestAIContextProvider(params string[] stateKeys) : AIContextProvider
+ {
+ public override IReadOnlyList StateKeys => stateKeys;
protected override ValueTask InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(context.AIContext);
@@ -1973,7 +2084,20 @@ public partial class ChatClientAgentTests
private sealed class TestChatHistoryProvider(string stateKey) : ChatHistoryProvider
{
- public override string StateKey => stateKey;
+ private readonly IReadOnlyList _stateKeys = [stateKey];
+
+ public override IReadOnlyList StateKeys => this._stateKeys;
+
+ protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
+ => new(context.RequestMessages);
+
+ protected override ValueTask InvokedCoreAsync(InvokedContext context, CancellationToken cancellationToken = default)
+ => default;
+ }
+
+ private sealed class MultiKeyTestChatHistoryProvider(params string[] stateKeys) : ChatHistoryProvider
+ {
+ public override IReadOnlyList StateKeys => stateKeys;
protected override ValueTask> InvokingCoreAsync(InvokingContext context, CancellationToken cancellationToken = default)
=> new(context.RequestMessages);
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs
index ebb1791dfd..1177a3c82a 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_BackgroundResponsesTests.cs
@@ -339,7 +339,7 @@ public class ChatClientAgent_BackgroundResponsesTests
// Create a mock chat history provider that would normally provide messages
var mockChatHistoryProvider = new Mock(null, null, null);
- mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+ mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
mockChatHistoryProvider
.Protected()
.Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -347,7 +347,7 @@ public class ChatClientAgent_BackgroundResponsesTests
// Create a mock AI context provider that would normally provide context
var mockContextProvider = new Mock(null, null, null);
- mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockContextProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -408,7 +408,7 @@ public class ChatClientAgent_BackgroundResponsesTests
// Create a mock chat history provider that would normally provide messages
var mockChatHistoryProvider = new Mock(null, null, null);
- mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+ mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
mockChatHistoryProvider
.Protected()
.Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -416,7 +416,7 @@ public class ChatClientAgent_BackgroundResponsesTests
// Create a mock AI context provider that would normally provide context
var mockContextProvider = new Mock(null, null, null);
- mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockContextProvider
.Protected()
.Setup>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -639,7 +639,7 @@ public class ChatClientAgent_BackgroundResponsesTests
List capturedMessagesAddedToProvider = [];
var mockChatHistoryProvider = new Mock(null, null, null);
- mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+ mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
mockChatHistoryProvider
.Protected()
.Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -648,7 +648,7 @@ public class ChatClientAgent_BackgroundResponsesTests
AIContextProvider.InvokedContext? capturedInvokedContext = null;
var mockContextProvider = new Mock(null, null, null);
- mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockContextProvider
.Protected()
.Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -703,7 +703,7 @@ public class ChatClientAgent_BackgroundResponsesTests
List capturedMessagesAddedToProvider = [];
var mockChatHistoryProvider = new Mock(null, null, null);
- mockChatHistoryProvider.SetupGet(p => p.StateKey).Returns("ChatHistoryProvider");
+ mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["ChatHistoryProvider"]);
mockChatHistoryProvider
.Protected()
.Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -712,7 +712,7 @@ public class ChatClientAgent_BackgroundResponsesTests
AIContextProvider.InvokedContext? capturedInvokedContext = null;
var mockContextProvider = new Mock(null, null, null);
- mockContextProvider.SetupGet(p => p.StateKey).Returns("Provider1");
+ mockContextProvider.SetupGet(p => p.StateKeys).Returns(["Provider1"]);
mockContextProvider
.Protected()
.Setup("InvokedCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs
index 59062cf49f..cc9b7acb19 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/ChatClient/ChatClientAgent_ChatHistoryManagementTests.cs
@@ -186,6 +186,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
It.IsAny())).ReturnsAsync(new ChatResponse([new(ChatRole.Assistant, "response")]));
Mock mockChatHistoryProvider = new(null, null, null);
+ mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
mockChatHistoryProvider
.Protected()
.Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -241,6 +242,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
It.IsAny())).Throws(new InvalidOperationException("Test Error"));
Mock mockChatHistoryProvider = new(null, null, null);
+ mockChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
mockChatHistoryProvider
.Protected()
.Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -430,6 +432,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
// Arrange a chat history provider to override the factory provided one.
Mock mockOverrideChatHistoryProvider = new(null, null, null);
+ mockOverrideChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
mockOverrideChatHistoryProvider
.Protected()
.Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
@@ -443,6 +446,7 @@ public class ChatClientAgent_ChatHistoryManagementTests
// Arrange a chat history provider to provide to the agent at construction time.
// This one shouldn't be used since it is being overridden.
Mock mockAgentOptionsChatHistoryProvider = new(null, null, null);
+ mockAgentOptionsChatHistoryProvider.SetupGet(p => p.StateKeys).Returns(["TestChatHistoryProvider"]);
mockAgentOptionsChatHistoryProvider
.Protected()
.Setup>>("InvokingCoreAsync", ItExpr.IsAny(), ItExpr.IsAny())
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs
index a0d6bbb35f..a782993f6a 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Data/TextSearchProviderTests.cs
@@ -39,17 +39,18 @@ public sealed class TextSearchProviderTests
}
[Fact]
- public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+ public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
{
// Arrange & Act
var provider = new TextSearchProvider((_, _) => Task.FromResult>([]));
// Assert
- Assert.Equal("TextSearchProvider", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("TextSearchProvider", provider.StateKeys);
}
[Fact]
- public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+ public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
{
// Arrange & Act
var provider = new TextSearchProvider(
@@ -57,7 +58,8 @@ public sealed class TextSearchProviderTests
new TextSearchProviderOptions { StateKey = "custom-key" });
// Assert
- Assert.Equal("custom-key", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("custom-key", provider.StateKeys);
}
[Theory]
diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs
index a0657d5a47..5211fa0956 100644
--- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/Memory/ChatHistoryMemoryProviderTests.cs
@@ -56,7 +56,7 @@ public class ChatHistoryMemoryProviderTests
}
[Fact]
- public void StateKey_ReturnsDefaultKey_WhenNoOptionsProvided()
+ public void StateKeys_ReturnsDefaultKey_WhenNoOptionsProvided()
{
// Arrange & Act
var provider = new ChatHistoryMemoryProvider(
@@ -66,11 +66,12 @@ public class ChatHistoryMemoryProviderTests
_ => new ChatHistoryMemoryProvider.State(new ChatHistoryMemoryProviderScope { UserId = "UID" }));
// Assert
- Assert.Equal("ChatHistoryMemoryProvider", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("ChatHistoryMemoryProvider", provider.StateKeys);
}
[Fact]
- public void StateKey_ReturnsCustomKey_WhenSetViaOptions()
+ public void StateKeys_ReturnsCustomKey_WhenSetViaOptions()
{
// Arrange & Act
var provider = new ChatHistoryMemoryProvider(
@@ -81,7 +82,8 @@ public class ChatHistoryMemoryProviderTests
new ChatHistoryMemoryProviderOptions { StateKey = "custom-key" });
// Assert
- Assert.Equal("custom-key", provider.StateKey);
+ Assert.Single(provider.StateKeys);
+ Assert.Contains("custom-key", provider.StateKeys);
}
[Fact]
diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeCodeGenTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeCodeGenTest.cs
index 459144a514..0efb0c19c4 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeCodeGenTest.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeCodeGenTest.cs
@@ -14,7 +14,7 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests;
public sealed class DeclarativeCodeGenTest(ITestOutputHelper output) : WorkflowTest(output)
{
[Theory]
- [InlineData("CheckSystem.yaml", "CheckSystem.json")]
+ [InlineData("CheckSystem.yaml", "CheckSystem.json", Skip = "Temporarily skipped")]
[InlineData("SendActivity.yaml", "SendActivity.json")]
[InlineData("InvokeAgent.yaml", "InvokeAgent.json")]
[InlineData("InvokeAgent.yaml", "InvokeAgent.json", true)]
diff --git a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeWorkflowTest.cs b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeWorkflowTest.cs
index c6a328d73d..eb1d0f55a2 100644
--- a/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeWorkflowTest.cs
+++ b/dotnet/tests/Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests/DeclarativeWorkflowTest.cs
@@ -15,7 +15,7 @@ namespace Microsoft.Agents.AI.Workflows.Declarative.IntegrationTests;
public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : WorkflowTest(output)
{
[Theory]
- [InlineData("CheckSystem.yaml", "CheckSystem.json")]
+ [InlineData("CheckSystem.yaml", "CheckSystem.json", Skip = "Temporarily skipped")]
[InlineData("ConversationMessages.yaml", "ConversationMessages.json")]
[InlineData("ConversationMessages.yaml", "ConversationMessages.json", true)]
[InlineData("InputArguments.yaml", "InputArguments.json")]
@@ -33,7 +33,7 @@ public sealed class DeclarativeWorkflowTest(ITestOutputHelper output) : Workflow
public Task ValidateScenarioAsync(string workflowFileName, string testcaseFileName, bool externalConveration = false) =>
this.RunWorkflowAsync(GetWorkflowPath(workflowFileName, isSample: true), testcaseFileName, externalConveration);
- [Theory]
+ [Theory(Skip = "Multi-turn tests hang in CI - needs investigation")]
[InlineData("ConfirmInput.yaml", "ConfirmInput.json", false)]
[InlineData("RequestExternalInput.yaml", "RequestExternalInput.json", false)]
public Task ValidateMultiTurnAsync(string workflowFileName, string testcaseFileName, bool isSample) =>
diff --git a/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs b/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs
index caa42ecc8d..e3b45bd5d2 100644
--- a/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs
+++ b/dotnet/tests/OpenAIAssistant.IntegrationTests/OpenAIAssistantStructuredOutputRunTests.cs
@@ -1,9 +1,23 @@
// Copyright (c) Microsoft. All rights reserved.
+using System.Threading.Tasks;
using AgentConformance.IntegrationTests;
namespace OpenAIAssistant.IntegrationTests;
public class OpenAIAssistantStructuredOutputRunTests() : StructuredOutputRunTests(() => new())
{
+ private const string SkipReason = "Fails intermittently on the build agent/CI";
+
+ [Fact(Skip = SkipReason)]
+ public override Task RunWithResponseFormatReturnsExpectedResultAsync() =>
+ base.RunWithResponseFormatReturnsExpectedResultAsync();
+
+ [Fact(Skip = SkipReason)]
+ public override Task RunWithGenericTypeReturnsExpectedResultAsync() =>
+ base.RunWithGenericTypeReturnsExpectedResultAsync();
+
+ [Fact(Skip = SkipReason)]
+ public override Task RunWithPrimitiveTypeReturnsExpectedResultAsync() =>
+ base.RunWithPrimitiveTypeReturnsExpectedResultAsync();
}
diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
index 997c375ed1..d8cf236add 100644
--- a/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
+++ b/python/packages/ag-ui/agent_framework_ag_ui/_run_common.py
@@ -372,6 +372,15 @@ def _emit_usage(content: Content) -> list[BaseEvent]:
return [CustomEvent(name="usage", value=usage_details)]
+def _emit_oauth_consent(content: Content) -> list[BaseEvent]:
+ """Emit an OAuth consent request as a custom event so frontends can render a consent link."""
+ return (
+ [CustomEvent(name="oauth_consent_request", value={"consent_link": content.consent_link})]
+ if content.consent_link
+ else []
+ )
+
+
def _emit_content(
content: Any,
flow: FlowState,
@@ -391,5 +400,7 @@ def _emit_content(
return _emit_approval_request(content, flow, predictive_handler, require_confirmation)
if content_type == "usage":
return _emit_usage(content)
+ if content_type == "oauth_consent_request":
+ return _emit_oauth_consent(content)
logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
return []
diff --git a/python/packages/ag-ui/tests/ag_ui/test_run.py b/python/packages/ag-ui/tests/ag_ui/test_run.py
index 73c9648c02..e0771e1b7e 100644
--- a/python/packages/ag-ui/tests/ag_ui/test_run.py
+++ b/python/packages/ag-ui/tests/ag_ui/test_run.py
@@ -4,6 +4,7 @@
import pytest
from ag_ui.core import (
+ CustomEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
@@ -871,3 +872,26 @@ class TestTextMessageEventBalancing:
assert len(start_events) == 2
assert len(end_events) == 2
+
+
+def test_emit_oauth_consent_request():
+ """Test that oauth_consent_request content emits a CustomEvent."""
+ content = Content.from_oauth_consent_request(
+ consent_link="https://login.microsoftonline.com/consent",
+ )
+ flow = FlowState()
+ events = _emit_content(content, flow)
+
+ assert len(events) == 1
+ assert isinstance(events[0], CustomEvent)
+ assert events[0].name == "oauth_consent_request"
+ assert events[0].value == {"consent_link": "https://login.microsoftonline.com/consent"}
+
+
+def test_emit_oauth_consent_request_no_link():
+ """Test that oauth_consent_request without a consent_link emits no events."""
+ content = Content("oauth_consent_request")
+ flow = FlowState()
+ events = _emit_content(content, flow)
+
+ assert len(events) == 0
diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py
index f9c2b99a6b..8ec2943181 100644
--- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py
+++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py
@@ -894,6 +894,7 @@ class AnthropicClient(
usage_details.append(Content.from_usage(usage_details=details))
return ChatResponseUpdate(
+ role="assistant",
response_id=event.message.id,
contents=[
*self._parse_contents_from_anthropic(event.message.content),
diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py
index 028e49673a..4f86c3eac2 100644
--- a/python/packages/anthropic/tests/test_anthropic_client.py
+++ b/python/packages/anthropic/tests/test_anthropic_client.py
@@ -1044,6 +1044,128 @@ async def test_inner_get_response_ignores_options_stream_streaming(mock_anthropi
assert mock_anthropic_client.beta.messages.create.call_args.kwargs["stream"] is True
+def test_process_stream_event_message_start_sets_assistant_role(mock_anthropic_client: MagicMock) -> None:
+ """Test that message_start streaming event sets role='assistant'.
+
+ This is critical: without role='assistant', _process_update cannot detect
+ a role boundary between a prior tool message and the new assistant turn,
+ causing tool_use blocks to collapse into a user-role message and triggering
+ Anthropic's '`tool_use` blocks can only be in `assistant` messages' error.
+ """
+ client = create_test_anthropic_client(mock_anthropic_client)
+
+ mock_event = MagicMock()
+ mock_event.type = "message_start"
+ mock_event.message.id = "msg_abc"
+ mock_event.message.role = "assistant"
+ mock_event.message.model = "claude-3-5-sonnet-20241022"
+ mock_event.message.content = []
+ mock_event.message.stop_reason = None
+ mock_event.message.usage = None
+
+ result = client._process_stream_event(mock_event)
+
+ assert result is not None
+ assert result.role == "assistant"
+
+
+def test_process_stream_event_message_start_role_prevents_tool_use_collapse() -> None:
+ """Regression test: tool_use blocks must not end up in a user-role message.
+
+ Simulates two consecutive streaming tool-call iterations:
+ Iteration 1: assistant emits tool_use → framework appends tool result (role=tool)
+ Iteration 2: assistant starts a new message_start → must create a NEW message
+
+ Without role='assistant' on the message_start update, _process_update sees
+ update.role=None (falsy) and appends to the last message (role='tool'),
+ producing {"role": "user", "content": [tool_result, tool_use]} which
+ Anthropic rejects with HTTP 400.
+ """
+ from agent_framework import ChatResponse, ChatResponseUpdate, Content, Message
+
+ # Simulate what the streaming tool loop produces after iteration 1:
+ # an existing 'tool' message is the last in the response
+ existing_tool_message = Message(
+ role="tool",
+ contents=[Content.from_function_result(call_id="call_1", result="some result")],
+ )
+
+ response = ChatResponse(messages=[existing_tool_message])
+
+ # Now simulate the message_start update from iteration 2 — WITH role set
+ message_start_update = ChatResponseUpdate(
+ role="assistant",
+ response_id="msg_iter2",
+ )
+
+ # Simulate a content_block_start carrying a tool_use — no role on this one (correct)
+ tool_use_update = ChatResponseUpdate(
+ contents=[
+ Content.from_function_call(
+ call_id="call_2",
+ name="get_weather",
+ arguments={"location": "NYC"},
+ )
+ ],
+ )
+
+ # Apply updates exactly as from_updates / _process_update would
+ from agent_framework._types import _process_update
+
+ _process_update(response, message_start_update)
+ _process_update(response, tool_use_update)
+
+ # Must have TWO messages: the original tool message + a new assistant message
+ assert len(response.messages) == 2, "tool_use from iteration 2 collapsed into the tool message from iteration 1"
+ assert response.messages[0].role == "tool"
+ assert response.messages[1].role == "assistant"
+
+ # The assistant message must contain the tool_use, not the tool result
+ assert response.messages[1].contents[0].type == "function_call"
+ assert response.messages[1].contents[0].call_id == "call_2"
+
+
+def test_process_stream_event_message_start_without_role_reproduces_bug() -> None:
+ """Documents the original bug: missing role causes tool_use to collapse into tool message.
+
+ This test demonstrates WHY the fix (adding role='assistant') was necessary.
+ It intentionally reproduces the broken behavior when role is absent.
+ """
+ from agent_framework import ChatResponse, ChatResponseUpdate, Content, Message
+ from agent_framework._types import _process_update
+
+ existing_tool_message = Message(
+ role="tool",
+ contents=[Content.from_function_result(call_id="call_1", result="some result")],
+ )
+ response = ChatResponse(messages=[existing_tool_message])
+
+ # message_start WITHOUT role (the original broken state)
+ message_start_update = ChatResponseUpdate(
+ role=None,
+ response_id="msg_iter2",
+ )
+ tool_use_update = ChatResponseUpdate(
+ contents=[
+ Content.from_function_call(
+ call_id="call_2",
+ name="get_weather",
+ arguments={"location": "NYC"},
+ )
+ ],
+ )
+
+ _process_update(response, message_start_update)
+ _process_update(response, tool_use_update)
+
+ # BUG: only 1 message — tool_use collapsed into the tool message
+ assert len(response.messages) == 1, "Expected bug: should still be 1 message without the fix"
+ # The single message has role='tool' but contains a function_call — invalid for Anthropic API
+ assert response.messages[0].role == "tool"
+ has_function_call = any(c.type == "function_call" for c in response.messages[0].contents)
+ assert has_function_call, "Expected bug: function_call leaked into tool message"
+
+
# Integration Tests
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
index 7590111bac..2c0498b1e4 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py
@@ -87,10 +87,11 @@ from azure.ai.agents.models import (
ToolApproval,
ToolDefinition,
ToolOutput,
+ VectorStoreDataSource,
)
from pydantic import BaseModel
-from ._shared import AzureAISettings, to_azure_ai_agent_tools
+from ._shared import AzureAISettings, resolve_file_ids, to_azure_ai_agent_tools
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -219,9 +220,21 @@ class AzureAIAgentClient(
# region Hosted Tool Factory Methods
@staticmethod
- def get_code_interpreter_tool() -> CodeInterpreterTool:
+ def get_code_interpreter_tool(
+ *,
+ file_ids: list[str | Content] | None = None,
+ data_sources: list[VectorStoreDataSource] | None = None,
+ ) -> CodeInterpreterTool:
"""Create a code interpreter tool configuration for Azure AI Agents.
+ Keyword Args:
+ file_ids: List of uploaded file IDs or Content objects to make available to
+ the code interpreter. Accepts plain strings or Content.from_hosted_file()
+ instances. The underlying SDK raises ValueError if both file_ids and
+ data_sources are provided.
+ data_sources: List of vector store data sources for enterprise file search.
+ Mutually exclusive with file_ids.
+
Returns:
A CodeInterpreterTool instance ready to pass to ChatAgent.
@@ -230,10 +243,21 @@ class AzureAIAgentClient(
from agent_framework.azure import AzureAIAgentClient
+ # Basic code interpreter
tool = AzureAIAgentClient.get_code_interpreter_tool()
+
+ # With uploaded file IDs
+ tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc123"])
+
+ # With Content objects
+ from agent_framework import Content
+
+ tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=[Content.from_hosted_file("file-abc123")])
+
agent = ChatAgent(client, tools=[tool])
"""
- return CodeInterpreterTool()
+ resolved = resolve_file_ids(file_ids)
+ return CodeInterpreterTool(file_ids=resolved, data_sources=data_sources)
@staticmethod
def get_file_search_tool(
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py
index 7c698847cc..61c4a09e94 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py
@@ -37,12 +37,13 @@ from agent_framework.openai._responses_client import RawOpenAIResponsesClient
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
ApproximateLocation,
+ CodeInterpreterContainerAuto,
CodeInterpreterTool,
- CodeInterpreterToolAuto,
+ FoundryFeaturesOptInKeys,
ImageGenTool,
MCPTool,
PromptAgentDefinition,
- PromptAgentDefinitionText,
+ PromptAgentDefinitionTextOptions,
RaiConfig,
Reasoning,
WebSearchPreviewTool,
@@ -50,7 +51,7 @@ from azure.ai.projects.models import (
from azure.ai.projects.models import FileSearchTool as ProjectsFileSearchTool
from azure.core.exceptions import ResourceNotFoundError
-from ._shared import AzureAISettings, create_text_format_config
+from ._shared import AzureAISettings, create_text_format_config, resolve_file_ids
if sys.version_info >= (3, 13):
from typing import TypeVar # type: ignore # pragma: no cover
@@ -78,6 +79,9 @@ class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False):
reasoning: Reasoning # type: ignore[misc]
"""Configuration for enabling reasoning capabilities (requires azure.ai.projects.models.Reasoning)."""
+ foundry_features: FoundryFeaturesOptInKeys | str
+ """Optional Foundry preview feature opt-in for agent version creation."""
+
AzureAIClientOptionsT = TypeVar(
"AzureAIClientOptionsT",
@@ -392,7 +396,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
# response_format is accessed from chat_options or additional_properties
# since the base class excludes it from run_options
if chat_options and (response_format := chat_options.get("response_format")):
- args["text"] = PromptAgentDefinitionText(format=create_text_format_config(response_format))
+ args["text"] = PromptAgentDefinitionTextOptions(format=create_text_format_config(response_format))
# Combine instructions from messages and options
# instructions is accessed from chat_options since the base class excludes it from run_options
@@ -404,11 +408,15 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
if combined_instructions:
args["instructions"] = "".join(combined_instructions)
- created_agent = await self.project_client.agents.create_version(
- agent_name=self.agent_name,
- definition=PromptAgentDefinition(**args),
- description=self.agent_description,
- )
+ create_version_kwargs: dict[str, Any] = {
+ "agent_name": self.agent_name,
+ "definition": PromptAgentDefinition(**args),
+ "description": self.agent_description,
+ }
+ if foundry_features := run_options.get("foundry_features"):
+ create_version_kwargs["foundry_features"] = foundry_features
+
+ created_agent = await self.project_client.agents.create_version(**create_version_kwargs)
self.agent_version = created_agent.version
self.warn_runtime_tools_and_structure_changed = True
@@ -500,6 +508,7 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
"temperature": ("temperature",),
"top_p": ("top_p",),
"reasoning": ("reasoning",),
+ "foundry_features": ("foundry_features",),
}
for run_keys in agent_level_option_to_run_keys.values():
@@ -526,9 +535,9 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"]))
if not self._is_application_endpoint:
- # Application-scoped response APIs do not support "agent" property.
+ # Application-scoped response APIs do not support "agent_reference" property.
agent_reference = await self._get_agent_reference_or_create(run_options, instructions, options)
- run_options["extra_body"] = {"agent": agent_reference}
+ run_options["extra_body"] = {"agent_reference": agent_reference}
# Remove only keys that map to this client's declared options TypedDict.
self._remove_agent_level_run_options(run_options, options)
@@ -588,6 +597,68 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
"""Get the current conversation ID from chat options or kwargs."""
return options.get("conversation_id") or kwargs.get("conversation_id") or self.conversation_id
+ @override
+ def _parse_response_from_openai(
+ self,
+ response: Any,
+ options: dict[str, Any],
+ ) -> ChatResponse:
+ """Parse an Azure AI Responses API response, handling Azure-specific output item types."""
+ result = super()._parse_response_from_openai(response, options)
+
+ if result.messages:
+ for item in response.output:
+ if item.type == "oauth_consent_request":
+ consent_link = item.consent_link
+ if consent_link and not consent_link.startswith("https://"):
+ logger.warning("Skipping oauth_consent_request with non-HTTPS consent_link: %s", item)
+ consent_link = ""
+ if consent_link:
+ result.messages[0].contents.append(
+ Content.from_oauth_consent_request(
+ consent_link=consent_link,
+ raw_representation=item,
+ )
+ )
+ else:
+ logger.warning("Received oauth_consent_request output without consent_link: %s", item)
+
+ return result
+
+ @override
+ def _parse_chunk_from_openai(
+ self,
+ event: Any,
+ options: dict[str, Any],
+ function_call_ids: dict[int, tuple[str, str]],
+ ) -> ChatResponseUpdate:
+ """Parse an Azure AI streaming event, handling Azure-specific event types."""
+ # Intercept output_item.added events for Azure-specific item types
+ if event.type == "response.output_item.added" and event.item.type == "oauth_consent_request":
+ event_item = event.item
+ consent_link = event_item.consent_link
+ if consent_link and not consent_link.startswith("https://"):
+ logger.warning("Skipping oauth_consent_request with non-HTTPS consent_link: %s", event_item)
+ consent_link = ""
+ contents: list[Content] = []
+ if consent_link:
+ contents.append(
+ Content.from_oauth_consent_request(
+ consent_link=consent_link,
+ raw_representation=event_item,
+ )
+ )
+ else:
+ logger.warning("Received oauth_consent_request output without consent_link: %s", event_item)
+ return ChatResponseUpdate(
+ contents=contents,
+ role="assistant",
+ model_id=self.model_id,
+ raw_representation=event,
+ )
+
+ return super()._parse_chunk_from_openai(event, options, function_call_ids)
+
def _prepare_messages_for_azure_ai(self, messages: Sequence[Message]) -> tuple[list[Message], str | None]:
"""Prepare input from messages and convert system/developer messages to instructions."""
result: list[Message] = []
@@ -830,14 +901,16 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
@staticmethod
def get_code_interpreter_tool( # type: ignore[override]
*,
- file_ids: list[str] | None = None,
+ file_ids: list[str | Content] | None = None,
container: Literal["auto"] | dict[str, Any] = "auto",
**kwargs: Any,
) -> CodeInterpreterTool:
"""Create a code interpreter tool configuration for Azure AI Projects.
Keyword Args:
- file_ids: Optional list of file IDs to make available to the code interpreter.
+ file_ids: Optional list of file IDs or Content objects to make available to
+ the code interpreter. Accepts plain strings or Content.from_hosted_file()
+ instances.
container: Container configuration. Use "auto" for automatic container management.
Note: Custom container settings from this parameter are not used by Azure AI Projects;
use file_ids instead.
@@ -857,7 +930,8 @@ class RawAzureAIClient(RawOpenAIResponsesClient[AzureAIClientOptionsT], Generic[
# Extract file_ids from container if provided as dict and file_ids not explicitly set
if file_ids is None and isinstance(container, dict):
file_ids = container.get("file_ids")
- tool_container = CodeInterpreterToolAuto(file_ids=file_ids if file_ids else None)
+ resolved = resolve_file_ids(file_ids)
+ tool_container = CodeInterpreterContainerAuto(file_ids=resolved)
return CodeInterpreterTool(container=tool_container, **kwargs)
@staticmethod
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
index eba210ff10..d02eb31bb6 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
@@ -18,7 +18,6 @@ from agent_framework._sessions import AgentSession, BaseContextProvider, Session
from agent_framework._settings import load_settings
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import ItemParam, ResponsesAssistantMessageItemParam, ResponsesUserMessageItemParam
from ._shared import AzureAISettings
@@ -149,7 +148,7 @@ class FoundryMemoryProvider(BaseContextProvider):
# On first run, retrieve static memories (user profile memories)
if not state.get("initialized"):
try:
- static_search_result = await self.project_client.memory_stores.search_memories(
+ static_search_result = await self.project_client.beta.memory_stores.search_memories(
name=self.memory_store_name,
scope=self.scope or context.session_id, # type: ignore[arg-type]
)
@@ -169,15 +168,15 @@ class FoundryMemoryProvider(BaseContextProvider):
if not has_input:
return
- # Convert input messages to ItemParam format for search
+ # Convert input messages to memory search item format
items = [
- ItemParam({"type": "text", "text": msg.text})
+ {"type": "text", "text": msg.text}
for msg in context.input_messages
if msg and msg.text and msg.text.strip()
]
try:
- search_result = await self.project_client.memory_stores.search_memories(
+ search_result = await self.project_client.beta.memory_stores.search_memories(
name=self.memory_store_name,
scope=self.scope or context.session_id, # type: ignore[arg-type]
items=items,
@@ -224,24 +223,24 @@ class FoundryMemoryProvider(BaseContextProvider):
if context.response and context.response.messages:
messages_to_store.extend(context.response.messages)
- # Filter and convert messages to ItemParam format
- items: list[ResponsesUserMessageItemParam | ResponsesAssistantMessageItemParam] = []
+ # Filter and convert messages to memory update item format
+ items: list[dict[str, str]] = []
for message in messages_to_store:
if message.role in {"user", "assistant", "system"} and message.text and message.text.strip():
if message.role == "user":
- items.append(ResponsesUserMessageItemParam(content=message.text))
+ items.append({"role": "user", "type": "message", "content": message.text})
elif message.role == "assistant":
- items.append(ResponsesAssistantMessageItemParam(content=message.text))
+ items.append({"role": "assistant", "type": "message", "content": message.text})
if not items:
return
try:
# Fire and forget - don't wait for the update to complete
- update_poller = await self.project_client.memory_stores.begin_update_memories(
+ update_poller = await self.project_client.beta.memory_stores.begin_update_memories(
name=self.memory_store_name,
scope=self.scope or context.session_id, # type: ignore[arg-type]
- items=items, # type: ignore[arg-type]
+ items=items,
previous_update_id=state.get("previous_update_id"),
update_delay=self.update_delay,
)
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
index 81276d446b..d6b922db91 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
@@ -4,7 +4,7 @@ from __future__ import annotations
import logging
import sys
-from collections.abc import Callable, MutableMapping, Sequence
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
from typing import Any, Generic
from agent_framework import (
@@ -21,10 +21,9 @@ from agent_framework._tools import ToolTypes
from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
- AgentReference,
AgentVersionDetails,
PromptAgentDefinition,
- PromptAgentDefinitionText,
+ PromptAgentDefinitionTextOptions,
)
from azure.ai.projects.models import (
FunctionTool as AzureFunctionTool,
@@ -200,13 +199,14 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
response_format = opts.get("response_format")
rai_config = opts.get("rai_config")
reasoning = opts.get("reasoning")
+ foundry_features = opts.get("foundry_features")
args: dict[str, Any] = {"model": resolved_model}
if instructions:
args["instructions"] = instructions
if response_format and isinstance(response_format, (type, dict)):
- args["text"] = PromptAgentDefinitionText(
+ args["text"] = PromptAgentDefinitionTextOptions(
format=create_text_format_config(response_format) # type: ignore[arg-type]
)
if rai_config:
@@ -241,11 +241,15 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
if all_tools_for_azure:
args["tools"] = to_azure_ai_tools(all_tools_for_azure)
- created_agent = await self._project_client.agents.create_version(
- agent_name=name,
- definition=PromptAgentDefinition(**args),
- description=description,
- )
+ create_version_kwargs: dict[str, Any] = {
+ "agent_name": name,
+ "definition": PromptAgentDefinition(**args),
+ "description": description,
+ }
+ if foundry_features:
+ create_version_kwargs["foundry_features"] = foundry_features
+
+ created_agent = await self._project_client.agents.create_version(**create_version_kwargs)
return self._to_chat_agent_from_details(
created_agent,
@@ -259,7 +263,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
self,
*,
name: str | None = None,
- reference: AgentReference | None = None,
+ reference: Mapping[str, str | None] | None = None,
tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
default_options: OptionsCoT | None = None,
middleware: Sequence[MiddlewareTypes] | None = None,
@@ -272,7 +276,7 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
Args:
name: The name of the agent to retrieve (fetches latest version).
- reference: Reference containing the agent's name and optionally a specific version.
+ reference: Mapping containing the agent's ``name`` and optionally a specific ``version``.
tools: Tools to make available to the agent. Required if the agent has function tools.
default_options: A TypedDict containing default chat options for the agent.
These options are applied to every run unless overridden.
@@ -287,12 +291,15 @@ class AzureAIProjectAgentProvider(Generic[OptionsCoT]):
"""
existing_agent: AgentVersionDetails
- if reference and reference.version:
+ reference_name = str(reference.get("name")) if reference and reference.get("name") else None
+ reference_version = str(reference.get("version")) if reference and reference.get("version") else None
+
+ if reference_name and reference_version:
# Fetch specific version
existing_agent = await self._project_client.agents.get_version(
- agent_name=reference.name, agent_version=reference.version
+ agent_name=reference_name, agent_version=reference_version
)
- elif agent_name := (reference.name if reference else name):
+ elif agent_name := (reference_name if reference_name else name):
# Fetch latest version
details = await self._project_client.agents.get(agent_name=agent_name)
existing_agent = details.versions.latest
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py
index 7dd1064bda..6f7d39c3be 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py
@@ -8,6 +8,7 @@ from collections.abc import Mapping, MutableMapping, Sequence
from typing import Any, cast
from agent_framework import (
+ Content,
FunctionTool,
)
from agent_framework.exceptions import IntegrationInvalidRequestException
@@ -18,9 +19,9 @@ from azure.ai.agents.models import (
from azure.ai.projects.models import (
CodeInterpreterTool,
MCPTool,
- ResponseTextFormatConfigurationJsonObject,
- ResponseTextFormatConfigurationJsonSchema,
- ResponseTextFormatConfigurationText,
+ TextResponseFormatConfigurationResponseFormatJsonObject,
+ TextResponseFormatConfigurationResponseFormatText,
+ TextResponseFormatJsonSchema,
Tool,
WebSearchPreviewTool,
)
@@ -109,6 +110,47 @@ def _extract_project_connection_id(additional_properties: dict[str, Any] | None)
return None
+def resolve_file_ids(file_ids: Sequence[str | Content] | None) -> list[str] | None:
+ """Resolve a list of file ID values that may include Content objects.
+
+ Accepts plain strings and Content objects with type "hosted_file", extracting
+ the file_id from each. This enables users to pass Content.from_hosted_file()
+ alongside plain file ID strings.
+
+ Args:
+ file_ids: Sequence of file ID strings or Content objects, or None.
+
+ Returns:
+ A list of resolved file ID strings, or None if input is None or empty.
+
+ Raises:
+ ValueError: If a Content object has an unsupported type (not "hosted_file").
+ """
+ if not file_ids:
+ return None
+
+ resolved: list[str] = []
+ for item in file_ids:
+ if isinstance(item, str):
+ if not item:
+ raise ValueError("file_ids must not contain empty strings.")
+ resolved.append(item)
+ elif isinstance(item, Content):
+ if item.type != "hosted_file":
+ raise ValueError(
+ f"Unsupported Content type '{item.type}' for code interpreter file_ids. "
+ "Only Content.from_hosted_file() is supported."
+ )
+ if item.file_id is None:
+ raise ValueError(
+ "Content.from_hosted_file() item is missing a file_id. "
+ "Ensure the Content object has a valid file_id before using it in file_ids."
+ )
+ resolved.append(item.file_id)
+
+ return resolved if resolved else None
+
+
def to_azure_ai_agent_tools(
tools: Sequence[FunctionTool | MutableMapping[str, Any]] | None,
run_options: dict[str, Any] | None = None,
@@ -421,9 +463,9 @@ def _prepare_mcp_tool_dict_for_azure_ai(tool_dict: dict[str, Any]) -> MCPTool:
def create_text_format_config(
response_format: type[BaseModel] | Mapping[str, Any],
) -> (
- ResponseTextFormatConfigurationJsonSchema
- | ResponseTextFormatConfigurationJsonObject
- | ResponseTextFormatConfigurationText
+ TextResponseFormatJsonSchema
+ | TextResponseFormatConfigurationResponseFormatJsonObject
+ | TextResponseFormatConfigurationResponseFormatText
):
"""Convert response_format into Azure text format configuration."""
if isinstance(response_format, type) and issubclass(response_format, BaseModel):
@@ -431,7 +473,7 @@ def create_text_format_config(
# Ensure additionalProperties is explicitly false to satisfy Azure validation
if isinstance(schema, dict):
schema.setdefault("additionalProperties", False)
- return ResponseTextFormatConfigurationJsonSchema(
+ return TextResponseFormatJsonSchema(
name=response_format.__name__,
schema=schema,
strict=True,
@@ -452,11 +494,11 @@ def create_text_format_config(
config_kwargs["strict"] = format_config["strict"]
if "description" in format_config:
config_kwargs["description"] = format_config["description"]
- return ResponseTextFormatConfigurationJsonSchema(**config_kwargs)
+ return TextResponseFormatJsonSchema(**config_kwargs)
if format_type == "json_object":
- return ResponseTextFormatConfigurationJsonObject()
+ return TextResponseFormatConfigurationResponseFormatJsonObject()
if format_type == "text":
- return ResponseTextFormatConfigurationText()
+ return TextResponseFormatConfigurationResponseFormatText()
raise IntegrationInvalidRequestException("response_format must be a Pydantic model or mapping.")
diff --git a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py
index b35efb6268..6c18352195 100644
--- a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py
+++ b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py
@@ -855,6 +855,110 @@ async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_file_search_with_
assert run_options["tool_resources"] == {"file_search": {"vector_store_ids": ["vs-123"]}}
+async def test_azure_ai_chat_client_prepare_tools_for_azure_ai_code_interpreter_with_file_ids(
+ mock_agents_client: MagicMock,
+) -> None:
+ """Test _prepare_tools_for_azure_ai with CodeInterpreterTool with file_ids from get_code_interpreter_tool()."""
+
+ client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent")
+
+ code_interpreter_tool = client.get_code_interpreter_tool(file_ids=["file-123", "file-456"])
+
+ run_options: dict[str, Any] = {}
+ result = await client._prepare_tools_for_azure_ai([code_interpreter_tool], run_options) # type: ignore
+
+ assert len(result) == 1
+ assert result[0] == {"type": "code_interpreter"}
+ assert "tool_resources" in run_options
+ assert "code_interpreter" in run_options["tool_resources"]
+ assert sorted(run_options["tool_resources"]["code_interpreter"]["file_ids"]) == ["file-123", "file-456"]
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_basic() -> None:
+ """Test get_code_interpreter_tool returns CodeInterpreterTool without files."""
+ from azure.ai.agents.models import CodeInterpreterTool
+
+ tool = AzureAIAgentClient.get_code_interpreter_tool()
+ assert isinstance(tool, CodeInterpreterTool)
+ assert len(tool.file_ids) == 0
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_file_ids() -> None:
+ """Test get_code_interpreter_tool forwards file_ids to the SDK."""
+ from azure.ai.agents.models import CodeInterpreterTool
+
+ tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc", "file-def"])
+ assert isinstance(tool, CodeInterpreterTool)
+ assert "file-abc" in tool.file_ids
+ assert "file-def" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_data_sources() -> None:
+ """Test get_code_interpreter_tool forwards data_sources to the SDK."""
+ from azure.ai.agents.models import CodeInterpreterTool, VectorStoreDataSource
+
+ ds = VectorStoreDataSource(asset_identifier="test-asset-id", asset_type="id_asset")
+ tool = AzureAIAgentClient.get_code_interpreter_tool(data_sources=[ds])
+ assert isinstance(tool, CodeInterpreterTool)
+ assert "test-asset-id" in tool.data_sources
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_mutually_exclusive() -> None:
+ """Test get_code_interpreter_tool raises ValueError when both file_ids and data_sources are provided."""
+ from azure.ai.agents.models import VectorStoreDataSource
+
+ ds = VectorStoreDataSource(asset_identifier="test-asset-id", asset_type="id_asset")
+ with pytest.raises(ValueError, match="mutually exclusive"):
+ AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-abc"], data_sources=[ds])
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_content() -> None:
+ """Test get_code_interpreter_tool accepts Content.from_hosted_file in file_ids."""
+ from agent_framework import Content
+ from azure.ai.agents.models import CodeInterpreterTool
+
+ content = Content.from_hosted_file("file-content-123")
+ tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
+ assert isinstance(tool, CodeInterpreterTool)
+ assert "file-content-123" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_with_mixed_file_ids() -> None:
+ """Test get_code_interpreter_tool accepts a mix of strings and Content objects."""
+ from agent_framework import Content
+ from azure.ai.agents.models import CodeInterpreterTool
+
+ content = Content.from_hosted_file("file-from-content")
+ tool = AzureAIAgentClient.get_code_interpreter_tool(file_ids=["file-plain", content])
+ assert isinstance(tool, CodeInterpreterTool)
+ assert "file-plain" in tool.file_ids
+ assert "file-from-content" in tool.file_ids
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_content_unsupported_type() -> None:
+ """Test get_code_interpreter_tool raises ValueError for unsupported Content types."""
+ from agent_framework import Content
+
+ content = Content.from_hosted_vector_store("vs-123")
+ with pytest.raises(ValueError, match="Unsupported Content type"):
+ AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_content_missing_file_id() -> None:
+ """Test get_code_interpreter_tool raises ValueError when Content.file_id is None."""
+ from agent_framework import Content
+
+ content = Content(type="hosted_file")
+ with pytest.raises(ValueError, match="missing a file_id"):
+ AzureAIAgentClient.get_code_interpreter_tool(file_ids=[content])
+
+
+async def test_azure_ai_chat_client_get_code_interpreter_tool_empty_string_file_id() -> None:
+ """Test get_code_interpreter_tool raises ValueError for empty string file_ids."""
+ with pytest.raises(ValueError, match="must not contain empty strings"):
+ AzureAIAgentClient.get_code_interpreter_tool(file_ids=[""])
+
+
async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals(
mock_agents_client: MagicMock,
) -> None:
diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py
index 4ec1b90971..e2145618c0 100644
--- a/python/packages/azure-ai/tests/test_azure_ai_client.py
+++ b/python/packages/azure-ai/tests/test_azure_ai_client.py
@@ -28,12 +28,12 @@ from agent_framework.openai._responses_client import RawOpenAIResponsesClient
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
ApproximateLocation,
+ CodeInterpreterContainerAuto,
CodeInterpreterTool,
- CodeInterpreterToolAuto,
FileSearchTool,
ImageGenTool,
MCPTool,
- ResponseTextFormatConfigurationJsonSchema,
+ TextResponseFormatJsonSchema,
WebSearchPreviewTool,
)
from azure.core.exceptions import ResourceNotFoundError
@@ -427,7 +427,7 @@ async def test_prepare_options_basic(mock_project_client: MagicMock) -> None:
run_options = await client._prepare_options(messages, {})
assert "extra_body" in run_options
- assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+ assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
@pytest.mark.parametrize(
@@ -465,7 +465,7 @@ async def test_prepare_options_with_application_endpoint(
if expects_agent:
assert "extra_body" in run_options
- assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+ assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
else:
assert "extra_body" not in run_options
@@ -507,7 +507,7 @@ async def test_prepare_options_with_application_project_client(
if expects_agent:
assert "extra_body" in run_options
- assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+ assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
else:
assert "extra_body" not in run_options
@@ -979,10 +979,10 @@ async def test_agent_creation_with_response_format(
assert hasattr(created_definition, "text")
assert created_definition.text is not None
- # Check that the format is a ResponseTextFormatConfigurationJsonSchema
+ # Check that the format is a TextResponseFormatJsonSchema
assert hasattr(created_definition.text, "format")
format_config = created_definition.text.format
- assert isinstance(format_config, ResponseTextFormatConfigurationJsonSchema)
+ assert isinstance(format_config, TextResponseFormatJsonSchema)
# Check the schema name matches the model class name
assert format_config.name == "ResponseFormatModel"
@@ -1040,7 +1040,7 @@ async def test_agent_creation_with_mapping_response_format(
assert hasattr(created_definition, "text")
assert created_definition.text is not None
format_config = created_definition.text.format
- assert isinstance(format_config, ResponseTextFormatConfigurationJsonSchema)
+ assert isinstance(format_config, TextResponseFormatJsonSchema)
assert format_config.name == runtime_schema["title"]
assert format_config.schema == runtime_schema
assert format_config.strict is True
@@ -1110,7 +1110,7 @@ async def test_prepare_options_excludes_response_format(
assert "text_format" not in run_options
# But extra_body should contain agent reference
assert "extra_body" in run_options
- assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+ assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
async def test_prepare_options_keeps_values_for_unsupported_option_keys(
@@ -1254,7 +1254,7 @@ def test_from_azure_ai_tools_mcp() -> None:
def test_from_azure_ai_tools_code_interpreter() -> None:
"""Test from_azure_ai_tools with Code Interpreter tool."""
- ci_tool = CodeInterpreterTool(container=CodeInterpreterToolAuto(file_ids=["file-1"]))
+ ci_tool = CodeInterpreterTool(container=CodeInterpreterContainerAuto(file_ids=["file-1"]))
parsed_tools = from_azure_ai_tools([ci_tool])
assert len(parsed_tools) == 1
assert parsed_tools[0]["type"] == "code_interpreter"
@@ -1685,6 +1685,35 @@ def test_get_code_interpreter_tool_with_file_ids() -> None:
assert tool["container"]["file_ids"] == ["file-123", "file-456"]
+def test_get_code_interpreter_tool_with_content() -> None:
+ """Test get_code_interpreter_tool accepts Content.from_hosted_file in file_ids."""
+ from agent_framework import Content
+
+ content = Content.from_hosted_file("file-content-123")
+ tool = AzureAIClient.get_code_interpreter_tool(file_ids=[content])
+ assert isinstance(tool, CodeInterpreterTool)
+ assert tool["container"]["file_ids"] == ["file-content-123"]
+
+
+def test_get_code_interpreter_tool_with_mixed_file_ids() -> None:
+ """Test get_code_interpreter_tool accepts a mix of strings and Content objects."""
+ from agent_framework import Content
+
+ content = Content.from_hosted_file("file-from-content")
+ tool = AzureAIClient.get_code_interpreter_tool(file_ids=["file-plain", content])
+ assert isinstance(tool, CodeInterpreterTool)
+ assert sorted(tool["container"]["file_ids"]) == ["file-from-content", "file-plain"]
+
+
+def test_get_code_interpreter_tool_content_unsupported_type() -> None:
+ """Test get_code_interpreter_tool raises ValueError for unsupported Content types."""
+ from agent_framework import Content
+
+ content = Content.from_hosted_vector_store("vs-123")
+ with pytest.raises(ValueError, match="Unsupported Content type"):
+ AzureAIClient.get_code_interpreter_tool(file_ids=[content])
+
+
def test_get_file_search_tool_basic() -> None:
"""Test get_file_search_tool returns FileSearchTool."""
tool = AzureAIClient.get_file_search_tool(vector_store_ids=["vs-123"])
@@ -2145,4 +2174,103 @@ def test_build_url_citation_content_with_dict(mock_project_client: MagicMock) ->
assert "get_url" not in ann.get("additional_properties", {})
+# region OAuth Consent
+
+
+def test_parse_chunk_with_oauth_consent_request(mock_project_client: MagicMock) -> None:
+ """Test that a streaming oauth_consent_request output item is parsed into oauth_consent_request content.
+
+ This reproduces the bug from issue #3950 where the event was logged as "Unparsed event"
+ and silently discarded, causing the agent run to complete with zero content.
+ """
+ client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+ chat_options: dict[str, Any] = {}
+ function_call_ids: dict[int, tuple[str, str]] = {}
+
+ mock_item = MagicMock()
+ mock_item.type = "oauth_consent_request"
+ mock_item.consent_link = "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc123"
+
+ mock_event = MagicMock()
+ mock_event.type = "response.output_item.added"
+ mock_event.item = mock_item
+ mock_event.output_index = 0
+
+ update = client._parse_chunk_from_openai(mock_event, chat_options, function_call_ids)
+
+ assert len(update.contents) == 1
+ consent_content = update.contents[0]
+ assert consent_content.type == "oauth_consent_request"
+ assert consent_content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc123"
+ assert consent_content.user_input_request is True
+
+
+def test_parse_response_with_oauth_consent_output_item(mock_project_client: MagicMock) -> None:
+ """Test that a non-streaming oauth_consent_request output item is parsed correctly."""
+ client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+
+ mock_item = MagicMock()
+ mock_item.type = "oauth_consent_request"
+ mock_item.consent_link = "https://login.microsoftonline.com/consent?code=abc"
+
+ mock_response = MagicMock()
+ mock_response.output = [mock_item]
+ mock_response.output_parsed = None
+ mock_response.metadata = {}
+ mock_response.id = "resp-oauth-1"
+ mock_response.model = "test-model"
+ mock_response.created_at = 1000000000
+ mock_response.usage = None
+ mock_response.status = "completed"
+
+ response = client._parse_response_from_openai(mock_response, {})
+
+ assert len(response.messages) > 0
+ consent_contents = [c for c in response.messages[0].contents if c.type == "oauth_consent_request"]
+ assert len(consent_contents) == 1
+ assert consent_contents[0].consent_link == "https://login.microsoftonline.com/consent?code=abc"
+
+
+def test_parse_chunk_oauth_consent_no_link(mock_project_client: MagicMock) -> None:
+ """Test that a streaming oauth_consent_request with no consent_link produces empty contents."""
+ client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+
+ mock_item = MagicMock()
+ mock_item.type = "oauth_consent_request"
+ mock_item.consent_link = ""
+
+ mock_event = MagicMock()
+ mock_event.type = "response.output_item.added"
+ mock_event.item = mock_item
+ mock_event.output_index = 0
+
+ update = client._parse_chunk_from_openai(mock_event, {}, {})
+
+ assert not any(c.type == "oauth_consent_request" for c in update.contents)
+
+
+def test_parse_response_oauth_consent_no_link(mock_project_client: MagicMock) -> None:
+ """Test that a non-streaming oauth_consent_request with no consent_link appends no content."""
+ client = AzureAIClient(project_client=mock_project_client, agent_name="test")
+
+ mock_item = MagicMock()
+ mock_item.type = "oauth_consent_request"
+ mock_item.consent_link = None
+
+ mock_response = MagicMock()
+ mock_response.output = [mock_item]
+ mock_response.output_parsed = None
+ mock_response.metadata = {}
+ mock_response.id = "resp-oauth-2"
+ mock_response.model = "test-model"
+ mock_response.created_at = 1000000000
+ mock_response.usage = None
+ mock_response.status = "completed"
+
+ response = client._parse_response_from_openai(mock_response, {})
+
+ consent_contents = [c for c in response.messages[0].contents if c.type == "oauth_consent_request"]
+ assert len(consent_contents) == 0
+
+
# endregion
diff --git a/python/packages/azure-ai/tests/test_foundry_memory_provider.py b/python/packages/azure-ai/tests/test_foundry_memory_provider.py
index 9c2968a65e..943a528968 100644
--- a/python/packages/azure-ai/tests/test_foundry_memory_provider.py
+++ b/python/packages/azure-ai/tests/test_foundry_memory_provider.py
@@ -17,9 +17,10 @@ from agent_framework_azure_ai._foundry_memory_provider import FoundryMemoryProvi
def mock_project_client() -> AsyncMock:
"""Create a mock AIProjectClient."""
mock_client = AsyncMock()
- mock_client.memory_stores = AsyncMock()
- mock_client.memory_stores.search_memories = AsyncMock()
- mock_client.memory_stores.begin_update_memories = AsyncMock()
+ mock_client.beta = AsyncMock()
+ mock_client.beta.memory_stores = AsyncMock()
+ mock_client.beta.memory_stores.search_memories = AsyncMock()
+ mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
return mock_client
@@ -146,7 +147,7 @@ class TestBeforeRun:
mem2.memory_item.content = "User is based in Seattle"
mock_search_result = Mock()
mock_search_result.memories = [mem1, mem2]
- mock_project_client.memory_stores.search_memories.return_value = mock_search_result
+ mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -161,7 +162,7 @@ class TestBeforeRun:
)
# Should call search_memories twice: once for static, once for contextual
- assert mock_project_client.memory_stores.search_memories.call_count == 2
+ assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Static memories should be cached
assert len(session.state[provider.source_id]["static_memories"]) == 2
assert session.state[provider.source_id]["initialized"] is True
@@ -181,7 +182,7 @@ class TestBeforeRun:
contextual_result.memories = [contextual_mem]
contextual_result.search_id = "search-123"
- mock_project_client.memory_stores.search_memories.side_effect = [static_result, contextual_result]
+ mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -208,7 +209,7 @@ class TestBeforeRun:
"""Empty input messages → only static search performed, no contextual search."""
static_result = Mock()
static_result.memories = []
- mock_project_client.memory_stores.search_memories.return_value = static_result
+ mock_project_client.beta.memory_stores.search_memories.return_value = static_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -223,14 +224,14 @@ class TestBeforeRun:
)
# Should only call search_memories once for static memories
- assert mock_project_client.memory_stores.search_memories.call_count == 1
+ assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
assert provider.source_id not in ctx.context_messages
async def test_empty_search_results_no_messages(self, mock_project_client: AsyncMock) -> None:
"""Empty search results → no messages added."""
mock_search_result = Mock()
mock_search_result.memories = []
- mock_project_client.memory_stores.search_memories.return_value = mock_search_result
+ mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -255,7 +256,7 @@ class TestBeforeRun:
contextual_result = Mock()
contextual_result.memories = []
- mock_project_client.memory_stores.search_memories.side_effect = [static_result, contextual_result]
+ mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -269,24 +270,24 @@ class TestBeforeRun:
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
- assert mock_project_client.memory_stores.search_memories.call_count == 2
+ assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
# Reset mock for second call
- mock_project_client.memory_stores.search_memories.reset_mock()
+ mock_project_client.beta.memory_stores.search_memories.reset_mock()
contextual_result2 = Mock()
contextual_result2.memories = []
- mock_project_client.memory_stores.search_memories.return_value = contextual_result2
+ mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
# Second call - should only search contextual, not static
ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
await provider.before_run( # type: ignore[arg-type]
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
- assert mock_project_client.memory_stores.search_memories.call_count == 1
+ assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
async def test_handles_search_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Search exception is logged but doesn't fail the operation."""
- mock_project_client.memory_stores.search_memories.side_effect = Exception("API error")
+ mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -315,7 +316,7 @@ class TestAfterRun:
"""Stores input+response messages via begin_update_memories."""
mock_poller = Mock()
mock_poller.update_id = "update-456"
- mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+ mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -330,8 +331,8 @@ class TestAfterRun:
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
- mock_project_client.memory_stores.begin_update_memories.assert_awaited_once()
- call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+ mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
+ call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["name"] == "test_store"
assert call_kwargs["scope"] == "user_123"
assert len(call_kwargs["items"]) == 2
@@ -342,7 +343,7 @@ class TestAfterRun:
async def test_only_stores_user_assistant_system(self, mock_project_client: AsyncMock) -> None:
"""Only stores user/assistant/system messages with text."""
mock_poller = Mock()
- mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+ mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -363,7 +364,7 @@ class TestAfterRun:
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
- call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+ call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
items = call_kwargs["items"]
assert len(items) == 2
assert items[0]["content"] == "hello"
@@ -390,12 +391,12 @@ class TestAfterRun:
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
- mock_project_client.memory_stores.begin_update_memories.assert_not_awaited()
+ mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock) -> None:
"""Uses the configured update_delay parameter."""
mock_poller = Mock()
- mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+ mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -411,7 +412,7 @@ class TestAfterRun:
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
)
- call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+ call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["update_delay"] == 60
async def test_uses_previous_update_id_for_incremental_updates(self, mock_project_client: AsyncMock) -> None:
@@ -421,7 +422,7 @@ class TestAfterRun:
mock_poller2 = Mock()
mock_poller2.update_id = "update-2"
- mock_project_client.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
+ mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
provider = FoundryMemoryProvider(
project_client=mock_project_client,
@@ -446,13 +447,13 @@ class TestAfterRun:
agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
)
- call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+ call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
assert call_kwargs["previous_update_id"] == "update-1"
assert session.state[provider.source_id]["previous_update_id"] == "update-2"
async def test_handles_update_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
"""Update exception is logged but doesn't fail the operation."""
- mock_project_client.memory_stores.begin_update_memories.side_effect = Exception("API error")
+ mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
provider = FoundryMemoryProvider(
project_client=mock_project_client,
diff --git a/python/packages/azure-ai/tests/test_provider.py b/python/packages/azure-ai/tests/test_provider.py
index 3765f17f1c..cb312983d4 100644
--- a/python/packages/azure-ai/tests/test_provider.py
+++ b/python/packages/azure-ai/tests/test_provider.py
@@ -8,7 +8,6 @@ from agent_framework import Agent, FunctionTool
from agent_framework._mcp import MCPTool
from azure.ai.projects.aio import AIProjectClient
from azure.ai.projects.models import (
- AgentReference,
AgentVersionDetails,
PromptAgentDefinition,
)
@@ -345,7 +344,7 @@ async def test_provider_get_agent_with_reference(mock_project_client: MagicMock)
mock_project_client.agents = AsyncMock()
mock_project_client.agents.get_version.return_value = mock_agent_version
- agent_reference = AgentReference(name="test-agent", version="1.0")
+ agent_reference = {"name": "test-agent", "version": "1.0"}
agent = await provider.get_agent(reference=agent_reference)
assert isinstance(agent, Agent)
diff --git a/python/packages/claude/agent_framework_claude/__init__.py b/python/packages/claude/agent_framework_claude/__init__.py
index 3c666f4a31..abf522fa4f 100644
--- a/python/packages/claude/agent_framework_claude/__init__.py
+++ b/python/packages/claude/agent_framework_claude/__init__.py
@@ -2,7 +2,7 @@
import importlib.metadata
-from ._agent import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings
+from ._agent import ClaudeAgent, ClaudeAgentOptions, ClaudeAgentSettings, RawClaudeAgent
try:
__version__ = importlib.metadata.version(__name__)
@@ -13,5 +13,6 @@ __all__ = [
"ClaudeAgent",
"ClaudeAgentOptions",
"ClaudeAgentSettings",
+ "RawClaudeAgent",
"__version__",
]
diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py
index 43f001b3db..f5aabc43a9 100644
--- a/python/packages/claude/agent_framework_claude/_agent.py
+++ b/python/packages/claude/agent_framework_claude/_agent.py
@@ -27,6 +27,7 @@ from agent_framework import (
normalize_tools,
)
from agent_framework.exceptions import AgentException
+from agent_framework.observability import AgentTelemetryLayer
from claude_agent_sdk import (
AssistantMessage,
ClaudeSDKClient,
@@ -171,8 +172,11 @@ OptionsT = TypeVar(
)
-class ClaudeAgent(BaseAgent, Generic[OptionsT]):
- """Claude Agent using Claude Code CLI.
+class RawClaudeAgent(BaseAgent, Generic[OptionsT]):
+ """Claude Agent using Claude Code CLI without telemetry layers.
+
+ This is the core Claude agent implementation without OpenTelemetry instrumentation.
+ For most use cases, prefer :class:`ClaudeAgent` which includes telemetry support.
Wraps the Claude Agent SDK to provide agentic capabilities including
tool use, session management, and streaming responses.
@@ -188,45 +192,13 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
.. code-block:: python
- from agent_framework_claude import ClaudeAgent
+ from agent_framework.anthropic import RawClaudeAgent
- async with ClaudeAgent(
+ async with RawClaudeAgent(
instructions="You are a helpful assistant.",
) as agent:
response = await agent.run("Hello!")
print(response.text)
-
- With streaming:
-
- .. code-block:: python
-
- async with ClaudeAgent() as agent:
- async for update in agent.run("Write a poem"):
- print(update.text, end="", flush=True)
-
- With session management:
-
- .. code-block:: python
-
- async with ClaudeAgent() as agent:
- session = agent.create_session()
- await agent.run("Remember my name is Alice", session=session)
- response = await agent.run("What's my name?", session=session)
- # Claude will remember "Alice" from the same session
-
- With Agent Framework tools:
-
- .. code-block:: python
-
- from agent_framework import tool
-
- @tool
- def greet(name: str) -> str:
- \"\"\"Greet someone by name.\"\"\"
- return f"Hello, {name}!"
-
- async with ClaudeAgent(tools=[greet]) as agent:
- response = await agent.run("Greet Alice")
"""
AGENT_PROVIDER_NAME: ClassVar[str] = "anthropic.claude"
@@ -246,7 +218,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
env_file_path: str | None = None,
env_file_encoding: str | None = None,
) -> None:
- """Initialize a ClaudeAgent instance.
+ """Initialize a RawClaudeAgent instance.
Args:
instructions: System prompt for the agent.
@@ -343,7 +315,7 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
normalized = normalize_tools(tool)
self._custom_tools.extend(normalized)
- async def __aenter__(self) -> ClaudeAgent[OptionsT]:
+ async def __aenter__(self) -> RawClaudeAgent[OptionsT]:
"""Start the agent when entering async context."""
await self.start()
return self
@@ -568,61 +540,19 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
return ""
return "\n".join([msg.text or "" for msg in messages])
- @overload
- def run(
- self,
- messages: AgentRunInputs | None = None,
- *,
- stream: Literal[True],
- session: AgentSession | None = None,
- options: OptionsT | MutableMapping[str, Any] | None = None,
- **kwargs: Any,
- ) -> AsyncIterable[AgentResponseUpdate]: ...
+ @property
+ def default_options(self) -> dict[str, Any]:
+ """Expose options with ``instructions`` key.
- @overload
- async def run(
- self,
- messages: AgentRunInputs | None = None,
- *,
- stream: Literal[False] = ...,
- session: AgentSession | None = None,
- options: OptionsT | MutableMapping[str, Any] | None = None,
- **kwargs: Any,
- ) -> AgentResponse[Any]: ...
-
- def run(
- self,
- messages: AgentRunInputs | None = None,
- *,
- stream: bool = False,
- session: AgentSession | None = None,
- options: OptionsT | MutableMapping[str, Any] | None = None,
- **kwargs: Any,
- ) -> AsyncIterable[AgentResponseUpdate] | Awaitable[AgentResponse[Any]]:
- """Run the agent with the given messages.
-
- Args:
- messages: The messages to process.
-
- Keyword Args:
- stream: If True, returns an async iterable of updates. If False (default),
- returns an awaitable AgentResponse.
- session: The conversation session. If session has service_session_id set,
- the agent will resume that session.
- options: Runtime options (model, permission_mode can be changed per-request).
- kwargs: Additional keyword arguments.
-
- Returns:
- When stream=True: An ResponseStream for streaming updates.
- When stream=False: An Awaitable[AgentResponse] with the complete response.
+ Maps ``system_prompt`` to ``instructions`` for compatibility with
+ :class:`AgentTelemetryLayer`, which reads the system prompt from
+ the ``instructions`` key.
"""
- response = ResponseStream(
- self._get_stream(messages, session=session, options=options, **kwargs),
- finalizer=self._finalize_response,
- )
- if stream:
- return response
- return response.get_final_response()
+ opts = dict(self._default_options)
+ system_prompt = opts.pop("system_prompt", None)
+ if system_prompt is not None:
+ opts["instructions"] = system_prompt
+ return opts
def _finalize_response(self, updates: Sequence[AgentResponseUpdate]) -> AgentResponse[Any]:
"""Build AgentResponse and propagate structured_output as value.
@@ -636,6 +566,61 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
structured_output = getattr(self, "_structured_output", None)
return AgentResponse.from_updates(updates, value=structured_output)
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = None,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = None,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
+ *,
+ stream: bool = False,
+ session: AgentSession | None = None,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
+ """Run the agent with the given messages.
+
+ Args:
+ messages: The messages to process.
+
+ Keyword Args:
+ stream: If True, returns an async iterable of updates. If False (default),
+ returns an awaitable AgentResponse.
+ session: The conversation session. If session has service_session_id set,
+ the agent will resume that session.
+ kwargs: Additional keyword arguments including 'options' for runtime options
+ (model, permission_mode can be changed per-request).
+
+ Returns:
+ When stream=True: An ResponseStream for streaming updates.
+ When stream=False: An Awaitable[AgentResponse] with the complete response.
+ """
+ options = kwargs.pop("options", None)
+ response = ResponseStream(
+ self._get_stream(messages, session=session, options=options, **kwargs),
+ finalizer=self._finalize_response,
+ )
+
+ if stream:
+ return response
+ return response.get_final_response()
+
async def _get_stream(
self,
messages: AgentRunInputs | None = None,
@@ -721,3 +706,25 @@ class ClaudeAgent(BaseAgent, Generic[OptionsT]):
# Store structured output for the finalizer
self._structured_output = structured_output
+
+
+class ClaudeAgent(AgentTelemetryLayer, RawClaudeAgent[OptionsT], Generic[OptionsT]):
+ """Claude Agent with OpenTelemetry instrumentation.
+
+ This is the recommended agent class for most use cases. It includes
+ OpenTelemetry-based telemetry for observability. For a minimal
+ implementation without telemetry, use :class:`RawClaudeAgent`.
+
+ Examples:
+ Basic usage with context manager:
+
+ .. code-block:: python
+
+ from agent_framework.anthropic import ClaudeAgent
+
+ async with ClaudeAgent(
+ instructions="You are a helpful assistant.",
+ ) as agent:
+ response = await agent.run("Hello!")
+ print(response.text)
+ """
diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py
index 0e126c36b9..e48a3b05d9 100644
--- a/python/packages/claude/tests/test_claude_agent.py
+++ b/python/packages/claude/tests/test_claude_agent.py
@@ -945,3 +945,191 @@ class TestClaudeAgentStructuredOutput:
with pytest.raises(AgentException) as exc_info:
await agent.run("Hello")
assert "Something went wrong" in str(exc_info.value)
+
+
+# region Test ClaudeAgent Telemetry
+
+
+class TestClaudeAgentTelemetry:
+ """Tests for ClaudeAgent OpenTelemetry instrumentation."""
+
+ @staticmethod
+ async def _create_async_generator(items: list[Any]) -> Any:
+ """Helper to create async generator from list."""
+ for item in items:
+ yield item
+
+ def _create_mock_client(self, messages: list[Any]) -> MagicMock:
+ """Create a mock ClaudeSDKClient that yields given messages."""
+ mock_client = MagicMock()
+ mock_client.connect = AsyncMock()
+ mock_client.disconnect = AsyncMock()
+ mock_client.query = AsyncMock()
+ mock_client.set_model = AsyncMock()
+ mock_client.set_permission_mode = AsyncMock()
+ mock_client.receive_response = MagicMock(return_value=self._create_async_generator(messages))
+ return mock_client
+
+ def _create_standard_messages(self) -> list[Any]:
+ """Create a standard set of mock messages for testing."""
+ from claude_agent_sdk import AssistantMessage, ResultMessage, TextBlock
+ from claude_agent_sdk.types import StreamEvent
+
+ return [
+ StreamEvent(
+ event={
+ "type": "content_block_delta",
+ "delta": {"type": "text_delta", "text": "Hello!"},
+ },
+ uuid="event-1",
+ session_id="session-123",
+ ),
+ AssistantMessage(
+ content=[TextBlock(text="Hello!")],
+ model="claude-sonnet",
+ ),
+ ResultMessage(
+ subtype="success",
+ duration_ms=100,
+ duration_api_ms=50,
+ is_error=False,
+ num_turns=1,
+ session_id="session-123",
+ ),
+ ]
+
+ async def test_run_emits_span_when_instrumentation_enabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Test that run() creates an OpenTelemetry span when instrumentation is enabled."""
+ from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+ messages = self._create_standard_messages()
+ mock_client = self._create_mock_client(messages)
+
+ monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+ with (
+ patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+ patch("agent_framework.observability._get_span") as mock_get_span,
+ ):
+ mock_span = MagicMock()
+ mock_get_span.return_value.__enter__ = MagicMock(return_value=mock_span)
+ mock_get_span.return_value.__exit__ = MagicMock(return_value=False)
+
+ agent = ClaudeAgent(name="test-agent")
+ response = await agent.run("Hello")
+
+ assert response.text == "Hello!"
+ mock_get_span.assert_called_once()
+ call_kwargs = mock_get_span.call_args[1]
+ assert call_kwargs["attributes"]["gen_ai.agent.name"] == "test-agent"
+ assert call_kwargs["attributes"]["gen_ai.operation.name"] == "invoke_agent"
+
+ async def test_run_skips_telemetry_when_instrumentation_disabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Test that run() skips telemetry when instrumentation is disabled."""
+ from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+ messages = self._create_standard_messages()
+ mock_client = self._create_mock_client(messages)
+
+ monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", False)
+
+ with (
+ patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+ patch("agent_framework.observability._get_span") as mock_get_span,
+ ):
+ agent = ClaudeAgent(name="test-agent")
+ response = await agent.run("Hello")
+
+ assert response.text == "Hello!"
+ mock_get_span.assert_not_called()
+
+ async def test_run_stream_emits_span_when_instrumentation_enabled(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Test that run(stream=True) creates a span when instrumentation is enabled."""
+ from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+ messages = self._create_standard_messages()
+ mock_client = self._create_mock_client(messages)
+
+ monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+ with (
+ patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+ patch("agent_framework.observability.get_tracer") as mock_get_tracer,
+ ):
+ mock_span = MagicMock()
+ mock_tracer = MagicMock()
+ mock_tracer.start_span.return_value = mock_span
+ mock_get_tracer.return_value = mock_tracer
+
+ agent = ClaudeAgent(name="stream-agent")
+ updates: list[AgentResponseUpdate] = []
+ async for update in agent.run("Hello", stream=True):
+ updates.append(update)
+
+ assert len(updates) == 1
+ mock_tracer.start_span.assert_called_once()
+ span_name = mock_tracer.start_span.call_args[0][0]
+ assert "stream-agent" in span_name
+ assert "invoke_agent" in span_name
+
+ async def test_run_captures_exception_in_span(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Test that exceptions during run() are captured in the telemetry span."""
+ from agent_framework.exceptions import AgentException
+ from agent_framework.observability import OBSERVABILITY_SETTINGS
+ from claude_agent_sdk import ResultMessage
+
+ error_messages = [
+ ResultMessage(
+ subtype="error",
+ duration_ms=100,
+ duration_api_ms=50,
+ is_error=True,
+ num_turns=0,
+ session_id="error-session",
+ result="Model not found",
+ ),
+ ]
+ mock_client = self._create_mock_client(error_messages)
+
+ monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+ with (
+ patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+ patch("agent_framework.observability._get_span") as mock_get_span,
+ patch("agent_framework.observability.capture_exception") as mock_capture_exc,
+ ):
+ mock_span = MagicMock()
+ mock_get_span.return_value.__enter__ = MagicMock(return_value=mock_span)
+ mock_get_span.return_value.__exit__ = MagicMock(return_value=False)
+
+ agent = ClaudeAgent(name="error-agent")
+ with pytest.raises(AgentException):
+ await agent.run("Hello")
+
+ mock_capture_exc.assert_called_once()
+ exc_kwargs = mock_capture_exc.call_args[1]
+ assert exc_kwargs["span"] is mock_span
+ assert isinstance(exc_kwargs["exception"], AgentException)
+
+ async def test_telemetry_uses_correct_provider_name(self, monkeypatch: pytest.MonkeyPatch) -> None:
+ """Test that telemetry uses AGENT_PROVIDER_NAME as provider."""
+ from agent_framework.observability import OBSERVABILITY_SETTINGS
+
+ messages = self._create_standard_messages()
+ mock_client = self._create_mock_client(messages)
+
+ monkeypatch.setattr(OBSERVABILITY_SETTINGS, "enable_instrumentation", True)
+
+ with (
+ patch("agent_framework_claude._agent.ClaudeSDKClient", return_value=mock_client),
+ patch("agent_framework.observability._get_span") as mock_get_span,
+ ):
+ mock_span = MagicMock()
+ mock_get_span.return_value.__enter__ = MagicMock(return_value=mock_span)
+ mock_get_span.return_value.__exit__ = MagicMock(return_value=False)
+
+ agent = ClaudeAgent(name="test-agent")
+ await agent.run("Hello")
+
+ call_kwargs = mock_get_span.call_args[1]
+ assert call_kwargs["attributes"]["gen_ai.provider.name"] == "anthropic.claude"
diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py
index 8f477f9223..cd2dc7bfc7 100644
--- a/python/packages/core/agent_framework/_agents.py
+++ b/python/packages/core/agent_framework/_agents.py
@@ -1051,10 +1051,11 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
else:
final_tools.append(tool) # type: ignore
+ existing_names = {name for t in final_tools if (name := _get_tool_name(t)) is not None}
for mcp_server in self.mcp_tools:
if not mcp_server.is_connected:
await self._async_exit_stack.enter_async_context(mcp_server)
- final_tools.extend(mcp_server.functions)
+ final_tools.extend(f for f in mcp_server.functions if f.name not in existing_names)
# Merge runtime kwargs into additional_function_arguments so they're available
# in function middleware context and tool invocation.
diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py
index beed97834c..ee0e813d27 100644
--- a/python/packages/core/agent_framework/_types.py
+++ b/python/packages/core/agent_framework/_types.py
@@ -345,6 +345,7 @@ ContentType = Literal[
"shell_command_output",
"function_approval_request",
"function_approval_response",
+ "oauth_consent_request",
]
@@ -498,6 +499,8 @@ class Content:
function_call: Content | None = None,
user_input_request: bool | None = None,
approved: bool | None = None,
+ # OAuth consent fields
+ consent_link: str | None = None,
# Common fields
annotations: Sequence[Annotation] | None = None,
additional_properties: MutableMapping[str, Any] | None = None,
@@ -546,6 +549,7 @@ class Content:
self.function_call = function_call
self.user_input_request = user_input_request
self.approved = approved
+ self.consent_link = consent_link
@classmethod
def from_text(
@@ -1122,6 +1126,37 @@ class Content:
raw_representation=raw_representation,
)
+ @classmethod
+ def from_oauth_consent_request(
+ cls: type[ContentT],
+ consent_link: str,
+ *,
+ annotations: Sequence[Annotation] | None = None,
+ additional_properties: MutableMapping[str, Any] | None = None,
+ raw_representation: Any = None,
+ ) -> ContentT:
+ """Create OAuth consent request content.
+
+ Args:
+ consent_link: The URL the user must visit to complete OAuth consent.
+
+ Keyword Args:
+ annotations: Optional annotations.
+ additional_properties: Optional additional properties.
+ raw_representation: Optional raw representation from the provider.
+
+ Returns:
+ A new Content instance with type ``oauth_consent_request``.
+ """
+ return cls(
+ "oauth_consent_request",
+ consent_link=consent_link,
+ user_input_request=True,
+ annotations=annotations,
+ additional_properties=additional_properties,
+ raw_representation=raw_representation,
+ )
+
def to_function_approval_response(
self,
approved: bool,
@@ -1176,6 +1211,7 @@ class Content:
"user_input_request",
"approved",
"id",
+ "consent_link",
"additional_properties",
)
diff --git a/python/packages/core/agent_framework/anthropic/__init__.py b/python/packages/core/agent_framework/anthropic/__init__.py
index 242554cf16..8be2a7d208 100644
--- a/python/packages/core/agent_framework/anthropic/__init__.py
+++ b/python/packages/core/agent_framework/anthropic/__init__.py
@@ -11,6 +11,7 @@ Supported classes:
- AnthropicChatOptions
- ClaudeAgent
- ClaudeAgentOptions
+- RawClaudeAgent
"""
import importlib
@@ -21,6 +22,7 @@ _IMPORTS: dict[str, tuple[str, str]] = {
"AnthropicChatOptions": ("agent_framework_anthropic", "agent-framework-anthropic"),
"ClaudeAgent": ("agent_framework_claude", "agent-framework-claude"),
"ClaudeAgentOptions": ("agent_framework_claude", "agent-framework-claude"),
+ "RawClaudeAgent": ("agent_framework_claude", "agent-framework-claude"),
}
diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml
index a3c3f53ed6..7f71f48de6 100644
--- a/python/packages/core/pyproject.toml
+++ b/python/packages/core/pyproject.toml
@@ -34,8 +34,7 @@ dependencies = [
# connectors and functions
"openai>=1.99.0",
"azure-identity>=1,<2",
- # Pinned to 2.0.0b3 - breaking changes in 2.0.0b4, unpin once upgrades complete
- "azure-ai-projects == 2.0.0b3",
+ "azure-ai-projects == 2.0.0b4",
"mcp[ws]>=1.24.0,<2",
"packaging>=24.1",
]
@@ -105,6 +104,7 @@ extend = "../../pyproject.toml"
[tool.pyright]
extends = "../../pyproject.toml"
+include = ["tests/workflow"]
[tool.mypy]
plugins = ['pydantic.mypy']
diff --git a/python/packages/core/tests/core/test_agents.py b/python/packages/core/tests/core/test_agents.py
index a857682fe2..c8d2d9bf8b 100644
--- a/python/packages/core/tests/core/test_agents.py
+++ b/python/packages/core/tests/core/test_agents.py
@@ -755,6 +755,49 @@ async def test_chat_agent_with_local_mcp_tools(client: SupportsChatGetResponse)
pass
+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]] = []
+
+ original_inner = chat_client_base._inner_get_response
+
+ async def capturing_inner(
+ *, messages: MutableSequence[Message], options: dict[str, Any], **kwargs: Any
+ ) -> ChatResponse:
+ captured_options.append(dict(options))
+ return await original_inner(messages=messages, options=options, **kwargs)
+
+ chat_client_base._inner_get_response = capturing_inner
+
+ # Create FunctionTool instances that simulate expanded MCP functions
+ mcp_func_a = FunctionTool(func=lambda: "a", name="tool_a", description="Tool A")
+ mcp_func_b = FunctionTool(func=lambda: "b", name="tool_b", description="Tool B")
+
+ # Create a mock MCP tool that is already connected (simulates turn 2)
+ mock_mcp_tool = MagicMock(spec=MCPTool)
+ mock_mcp_tool.is_connected = True
+ mock_mcp_tool.functions = [mcp_func_a, mcp_func_b]
+ mock_mcp_tool.__aenter__ = AsyncMock(return_value=mock_mcp_tool)
+ mock_mcp_tool.__aexit__ = AsyncMock(return_value=None)
+
+ # Agent has the MCP tool in its constructor (stored in self.mcp_tools)
+ agent = Agent(client=chat_client_base, name="TestAgent", tools=[mock_mcp_tool])
+
+ # Simulate AG-UI turn 2: pass already-expanded MCP functions + a client tool as runtime tools
+ client_tool = FunctionTool(func=lambda: "client", name="client_tool", description="Client tool")
+ runtime_tools = [mcp_func_a, mcp_func_b, client_tool]
+
+ await agent.run("hello", tools=runtime_tools)
+
+ # Verify the chat client received each tool exactly once
+ assert len(captured_options) >= 1
+ tool_names = [t.name for t in captured_options[0]["tools"]]
+ assert tool_names.count("tool_a") == 1, f"tool_a duplicated: {tool_names}"
+ assert tool_names.count("tool_b") == 1, f"tool_b duplicated: {tool_names}"
+ assert "client_tool" in tool_names
+ assert len(tool_names) == 3
+
+
async def test_agent_tool_receives_session_in_kwargs(chat_client_base: Any) -> None:
"""Verify tool execution receives 'session' inside **kwargs when function is called by client."""
diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py
index c858ff1e3f..bcf3a6891b 100644
--- a/python/packages/core/tests/core/test_types.py
+++ b/python/packages/core/tests/core/test_types.py
@@ -3424,3 +3424,30 @@ class TestResponseStreamEdgeCases:
# endregion
+
+
+# region OAuth Consent Content
+
+
+def test_oauth_consent_request_creation():
+ """Test Content.from_oauth_consent_request creates the correct content."""
+ content = Content.from_oauth_consent_request(
+ consent_link="https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc",
+ )
+ assert content.type == "oauth_consent_request"
+ assert content.consent_link == "https://login.microsoftonline.com/common/oauth2/authorize?client_id=abc"
+ assert content.user_input_request is True
+
+
+def test_oauth_consent_request_serialization_roundtrip():
+ """Test that oauth_consent_request content serializes and includes consent_link."""
+ content = Content.from_oauth_consent_request(
+ consent_link="https://login.microsoftonline.com/consent",
+ )
+ d = content.to_dict()
+ assert d["type"] == "oauth_consent_request"
+ assert d["consent_link"] == "https://login.microsoftonline.com/consent"
+ assert d["user_input_request"] is True
+
+
+# endregion
diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py
index 4a850db642..788e96e61e 100644
--- a/python/packages/core/tests/workflow/test_agent_executor.py
+++ b/python/packages/core/tests/workflow/test_agent_executor.py
@@ -2,19 +2,20 @@
import logging
from collections.abc import AsyncIterable, Awaitable
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal, overload
import pytest
-
from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
+ AgentRunInputs,
AgentSession,
BaseAgent,
Content,
Message,
ResponseStream,
+ WorkflowEvent,
WorkflowRunState,
)
from agent_framework._workflows._agent_executor import AgentExecutorResponse
@@ -32,26 +33,56 @@ class _CountingAgent(BaseAgent):
super().__init__(**kwargs)
self.call_count = 0
+ @overload
def run(
self,
- messages: str | Message | list[str] | list[Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> (
+ Awaitable[AgentResponse[Any]]
+ | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
+ ):
self.call_count += 1
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(
- contents=[Content.from_text(text=f"Response #{self.call_count}: {self.name}")]
+ contents=[
+ Content.from_text(
+ text=f"Response #{self.call_count}: {self.name}"
+ )
+ ]
)
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
async def _run() -> AgentResponse:
- return AgentResponse(messages=[Message("assistant", [f"Response #{self.call_count}: {self.name}"])])
+ return AgentResponse(
+ messages=[
+ Message("assistant", [f"Response #{self.call_count}: {self.name}"])
+ ]
+ )
return _run()
@@ -63,13 +94,36 @@ class _StreamingHookAgent(BaseAgent):
super().__init__(**kwargs)
self.result_hook_called = False
+ @overload
def run(
self,
- messages: str | Message | list[str] | list[Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
+ session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> (
+ Awaitable[AgentResponse[Any]]
+ | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]
+ ):
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -78,13 +132,15 @@ class _StreamingHookAgent(BaseAgent):
role="assistant",
)
- async def _mark_result_hook_called(response: AgentResponse) -> AgentResponse:
+ async def _mark_result_hook_called(
+ response: AgentResponse,
+ ) -> AgentResponse:
self.result_hook_called = True
return response
- return ResponseStream(_stream(), finalizer=AgentResponse.from_updates).with_result_hook(
- _mark_result_hook_called
- )
+ return ResponseStream(
+ _stream(), finalizer=AgentResponse.from_updates
+ ).with_result_hook(_mark_result_hook_called)
async def _run() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", ["hook test"])])
@@ -92,7 +148,9 @@ class _StreamingHookAgent(BaseAgent):
return _run()
-async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> None:
+async def test_agent_executor_streaming_finalizes_stream_and_runs_result_hooks() -> (
+ None
+):
"""AgentExecutor should call get_final_response() so stream result hooks execute."""
agent = _StreamingHookAgent(id="hook_agent", name="HookAgent")
executor = AgentExecutor(agent, id="hook_exec")
@@ -159,7 +217,9 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
executor_state = executor_states[executor.id] # type: ignore[index]
assert "cache" in executor_state, "Checkpoint should store executor cache state"
- assert "agent_session" in executor_state, "Checkpoint should store executor session state"
+ assert "agent_session" in executor_state, (
+ "Checkpoint should store executor session state"
+ )
# Verify session state structure
session_state = executor_state["agent_session"] # type: ignore[index]
@@ -180,11 +240,15 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
assert restored_agent.call_count == 0
# Build new workflow with the restored executor
- wf_resume = SequentialBuilder(participants=[restored_executor], checkpoint_storage=storage).build()
+ wf_resume = SequentialBuilder(
+ participants=[restored_executor], checkpoint_storage=storage
+ ).build()
# Resume from checkpoint
resumed_output: AgentExecutorResponse | None = None
- async for ev in wf_resume.run(checkpoint_id=restore_checkpoint.checkpoint_id, stream=True):
+ async for ev in wf_resume.run(
+ checkpoint_id=restore_checkpoint.checkpoint_id, stream=True
+ ):
if ev.type == "output":
resumed_output = ev.data # type: ignore[assignment]
if ev.type == "status" and ev.state in (
@@ -278,7 +342,7 @@ async def test_agent_executor_run_streaming_with_stream_kwarg_does_not_raise() -
workflow = SequentialBuilder(participants=[executor]).build()
# stream=True at workflow level triggers streaming mode (returns async iterable)
- events = []
+ events: list[WorkflowEvent] = []
async for event in workflow.run("hello", stream=True):
events.append(event)
assert len(events) > 0
@@ -288,10 +352,13 @@ async def test_agent_executor_run_streaming_with_stream_kwarg_does_not_raise() -
@pytest.mark.parametrize("reserved_kwarg", ["session", "stream", "messages"])
async def test_prepare_agent_run_args_strips_reserved_kwargs(reserved_kwarg: str, caplog: "LogCaptureFixture") -> None:
"""_prepare_agent_run_args must remove reserved kwargs and log a warning."""
- raw = {reserved_kwarg: "should-be-stripped", "custom_key": "keep-me"}
+ raw: dict[str, Any] = {
+ reserved_kwarg: "should-be-stripped",
+ "custom_key": "keep-me",
+ }
with caplog.at_level(logging.WARNING):
- run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)
+ run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw) # pyright: ignore[reportPrivateUsage]
assert reserved_kwarg not in run_kwargs
assert "custom_key" in run_kwargs
@@ -302,8 +369,8 @@ async def test_prepare_agent_run_args_strips_reserved_kwargs(reserved_kwarg: str
async def test_prepare_agent_run_args_preserves_non_reserved_kwargs() -> None:
"""Non-reserved workflow kwargs should pass through unchanged."""
- raw = {"custom_param": "value", "another": 42}
- run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)
+ raw: dict[str, Any] = {"custom_param": "value", "another": 42}
+ run_kwargs, _options = AgentExecutor._prepare_agent_run_args(raw) # pyright: ignore[reportPrivateUsage]
assert run_kwargs["custom_param"] == "value"
assert run_kwargs["another"] == 42
@@ -312,10 +379,10 @@ async def test_prepare_agent_run_args_strips_all_reserved_kwargs_at_once(
caplog: "LogCaptureFixture",
) -> None:
"""All reserved kwargs should be stripped when supplied together, each emitting a warning."""
- raw = {"session": "x", "stream": True, "messages": [], "custom": 1}
+ raw: dict[str, Any] = {"session": "x", "stream": True, "messages": [], "custom": 1}
with caplog.at_level(logging.WARNING):
- run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw)
+ run_kwargs, options = AgentExecutor._prepare_agent_run_args(raw) # pyright: ignore[reportPrivateUsage]
assert "session" not in run_kwargs
assert "stream" not in run_kwargs
@@ -324,7 +391,11 @@ async def test_prepare_agent_run_args_strips_all_reserved_kwargs_at_once(
assert options is not None
assert options["additional_function_arguments"]["custom"] == 1
- warned_keys = {r.message.split("'")[1] for r in caplog.records if "reserved" in r.message.lower()}
+ warned_keys = {
+ r.message.split("'")[1]
+ for r in caplog.records
+ if "reserved" in r.message.lower()
+ }
assert warned_keys == {"session", "stream", "messages"}
diff --git a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
index cae5ea4e3b..07a37f9617 100644
--- a/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
+++ b/python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
@@ -3,7 +3,7 @@
"""Tests for AgentExecutor handling of tool calls and results in streaming mode."""
from collections.abc import AsyncIterable, Awaitable, Mapping, Sequence
-from typing import Any
+from typing import Any, Literal, overload
from typing_extensions import Never
@@ -13,6 +13,7 @@ from agent_framework import (
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
+ AgentRunInputs,
AgentSession,
BaseAgent,
ChatResponse,
@@ -37,18 +38,38 @@ class _ToolCallingAgent(BaseAgent):
def __init__(self, **kwargs: Any) -> None:
super().__init__(**kwargs)
+ @overload
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
return ResponseStream(self._run_stream_impl(), finalizer=AgentResponse.from_updates)
- async def _run() -> AgentResponse:
+ async def _run() -> AgentResponse[Any]:
return AgentResponse(messages=[Message("assistant", ["done"])])
return _run()
@@ -111,6 +132,7 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
# First event: text update
assert events[0].data is not None
assert events[0].data.contents[0].type == "text"
+ assert events[0].data.contents[0].text is not None
assert "Let me search" in events[0].data.contents[0].text
# Second event: function call
@@ -129,6 +151,7 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
# Fourth event: final text
assert events[3].data is not None
assert events[3].data.contents[0].type == "text"
+ assert events[3].data.contents[0].text is not None
assert "sunny" in events[3].data.contents[0].text
diff --git a/python/packages/core/tests/workflow/test_agent_utils.py b/python/packages/core/tests/workflow/test_agent_utils.py
index d3889b4d3b..07d1e64c08 100644
--- a/python/packages/core/tests/workflow/test_agent_utils.py
+++ b/python/packages/core/tests/workflow/test_agent_utils.py
@@ -1,9 +1,9 @@
# Copyright (c) Microsoft. All rights reserved.
-from collections.abc import AsyncIterable
-from typing import Any
+from collections.abc import Awaitable
+from typing import Any, Literal, overload
-from agent_framework import AgentResponse, AgentResponseUpdate, AgentSession, Message
+from agent_framework import AgentResponse, AgentResponseUpdate, AgentRunInputs, AgentSession, ResponseStream
from agent_framework._workflows._agent_utils import resolve_agent_id
@@ -11,40 +11,23 @@ class MockAgent:
"""Mock agent for testing agent utilities."""
def __init__(self, agent_id: str, name: str | None = None) -> None:
- self._id = agent_id
- self._name = name
+ self.id: str = agent_id
+ self.name: str | None = name
+ self.description: str | None = None
- @property
- def id(self) -> str:
- return self._id
-
- @property
- def name(self) -> str | None:
- return self._name
-
- @property
- def display_name(self) -> str:
- """Returns the display name of the agent."""
- ...
-
- @property
- def description(self) -> str | None:
- """Returns the description of the agent."""
- ...
-
- def run(
- self,
- messages: str | Message | list[str] | list[Message] | None = None,
- *,
- stream: bool = False,
- session: AgentSession | None = None,
- **kwargs: Any,
- ) -> AgentResponse | AsyncIterable[AgentResponseUpdate]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+ def run(self, messages: AgentRunInputs | None = None, *, stream: bool = False, session: AgentSession | None = None, **kwargs: Any) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def create_session(self, **kwargs: Any) -> AgentSession:
"""Creates a new conversation session for the agent."""
...
+ def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
+ return AgentSession()
+
def test_resolve_agent_id_with_name() -> None:
"""Test that resolve_agent_id returns name when agent has a name."""
diff --git a/python/packages/core/tests/workflow/test_checkpoint.py b/python/packages/core/tests/workflow/test_checkpoint.py
index b05d625502..a32489acc0 100644
--- a/python/packages/core/tests/workflow/test_checkpoint.py
+++ b/python/packages/core/tests/workflow/test_checkpoint.py
@@ -5,6 +5,7 @@ import tempfile
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
+from typing import Any
import pytest
@@ -24,7 +25,7 @@ class _TestToolApprovalRequest:
"""Request data for tool approval in tests."""
tool_name: str
- arguments: dict
+ arguments: dict[str, Any]
timestamp: datetime
@@ -41,7 +42,7 @@ class _TestApprovalRequest:
"""Approval request data for tests."""
action: str
- params: tuple
+ params: tuple[Any, ...]
@dataclass
@@ -78,8 +79,8 @@ def test_workflow_checkpoint_custom_values():
workflow_name="test-workflow-456",
graph_signature_hash="test-hash-456",
timestamp=custom_timestamp,
- messages={"executor1": [{"data": "test"}]},
- pending_request_info_events={"req123": {"data": "test"}},
+ messages={"executor1": [{"data": "test"}]}, # type: ignore[arg-type] # raw dict for serialization test
+ pending_request_info_events={"req123": {"data": "test"}}, # type: ignore[arg-type] # raw dict for serialization test
state={"key": "value"},
iteration_count=5,
metadata={"test": True},
@@ -103,7 +104,7 @@ def test_workflow_checkpoint_to_dict():
checkpoint_id="test-id",
workflow_name="test-workflow",
graph_signature_hash="test-hash",
- messages={"executor1": [{"data": "test"}]},
+ messages={"executor1": [{"data": "test"}]}, # type: ignore[arg-type] # raw dict for serialization test
state={"key": "value"},
iteration_count=5,
)
@@ -161,8 +162,8 @@ async def test_memory_checkpoint_storage_save_and_load():
checkpoint = WorkflowCheckpoint(
workflow_name="test-workflow",
graph_signature_hash="test-hash",
- messages={"executor1": [{"data": "hello"}]},
- pending_request_info_events={"req123": {"data": "test"}},
+ messages={"executor1": [{"data": "hello"}]}, # type: ignore[arg-type] # raw dict for serialization test
+ pending_request_info_events={"req123": {"data": "test"}}, # type: ignore[arg-type] # raw dict for serialization test
)
# Save checkpoint
@@ -776,9 +777,9 @@ async def test_file_checkpoint_storage_save_and_load():
checkpoint = WorkflowCheckpoint(
workflow_name="test-workflow",
graph_signature_hash="test-hash",
- messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},
+ messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]}, # type: ignore[arg-type] # raw dict for serialization test
state={"key": "value"},
- pending_request_info_events={"req123": {"data": "test"}},
+ pending_request_info_events={"req123": {"data": "test"}}, # type: ignore[arg-type] # raw dict for serialization test
)
# Save checkpoint
@@ -904,9 +905,9 @@ async def test_file_checkpoint_storage_json_serialization():
checkpoint = WorkflowCheckpoint(
workflow_name="test-workflow",
graph_signature_hash="test-hash",
- messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
+ messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]}, # type: ignore[arg-type] # raw dict for serialization test
state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
- pending_request_info_events={"req123": {"data": "test"}},
+ pending_request_info_events={"req123": {"data": "test"}}, # type: ignore[arg-type] # raw dict for serialization test
)
# Save and load
diff --git a/python/packages/core/tests/workflow/test_checkpoint_encode.py b/python/packages/core/tests/workflow/test_checkpoint_encode.py
index 68ec1ac4e3..02da2f1297 100644
--- a/python/packages/core/tests/workflow/test_checkpoint_encode.py
+++ b/python/packages/core/tests/workflow/test_checkpoint_encode.py
@@ -3,11 +3,11 @@
import json
from dataclasses import dataclass
from datetime import datetime, timezone
-from typing import Any
+from typing import Any, cast
from agent_framework._workflows._checkpoint_encoding import (
- _PICKLE_MARKER,
- _TYPE_MARKER,
+ _PICKLE_MARKER, # pyright: ignore[reportPrivateUsage]
+ _TYPE_MARKER, # pyright: ignore[reportPrivateUsage]
encode_checkpoint_value,
)
@@ -185,8 +185,9 @@ def test_encode_list_of_dataclasses() -> None:
result = encode_checkpoint_value(data)
assert isinstance(result, list)
- assert len(result) == 2
- for item in result:
+ result_list = cast(list[Any], result)
+ assert len(result_list) == 2
+ for item in result_list:
assert _PICKLE_MARKER in item
diff --git a/python/packages/core/tests/workflow/test_edge.py b/python/packages/core/tests/workflow/test_edge.py
index f63cf9b45b..ecaa341726 100644
--- a/python/packages/core/tests/workflow/test_edge.py
+++ b/python/packages/core/tests/workflow/test_edge.py
@@ -4,6 +4,8 @@ from dataclasses import dataclass
from typing import Any
from unittest.mock import patch
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
import pytest
from agent_framework import (
@@ -275,6 +277,7 @@ async def test_single_edge_group_send_message_with_condition_pass() -> None:
success = await edge_runner.send_message(message, state, ctx)
assert success is True
assert target.call_count == 1
+ assert target.last_message is not None
assert target.last_message.data == "test"
@@ -301,7 +304,7 @@ async def test_single_edge_group_send_message_with_condition_fail() -> None:
assert target.call_count == 0
-async def test_single_edge_group_tracing_success(span_exporter) -> None:
+async def test_single_edge_group_tracing_success(span_exporter: InMemorySpanExporter) -> None:
"""Test that single edge group processing creates proper success spans."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
@@ -352,7 +355,7 @@ async def test_single_edge_group_tracing_success(span_exporter) -> None:
assert link.context.span_id == int("00f067aa0ba902b7", 16)
-async def test_single_edge_group_tracing_condition_failure(span_exporter) -> None:
+async def test_single_edge_group_tracing_condition_failure(span_exporter: InMemorySpanExporter) -> None:
"""Test that single edge group processing creates proper spans for condition failures."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
@@ -386,7 +389,7 @@ async def test_single_edge_group_tracing_condition_failure(span_exporter) -> Non
assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_CONDITION_FALSE.value
-async def test_single_edge_group_tracing_type_mismatch(span_exporter) -> None:
+async def test_single_edge_group_tracing_type_mismatch(span_exporter: InMemorySpanExporter) -> None:
"""Test that single edge group processing creates proper spans for type mismatches."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
@@ -421,7 +424,7 @@ async def test_single_edge_group_tracing_type_mismatch(span_exporter) -> None:
assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value
-async def test_single_edge_group_tracing_target_mismatch(span_exporter) -> None:
+async def test_single_edge_group_tracing_target_mismatch(span_exporter: InMemorySpanExporter) -> None:
"""Test that single edge group processing creates proper spans for target mismatches."""
source = MockExecutor(id="source_executor")
target = MockExecutor(id="target_executor")
@@ -775,7 +778,7 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_in
assert success is False
-async def test_fan_out_edge_group_tracing_success(span_exporter) -> None:
+async def test_fan_out_edge_group_tracing_success(span_exporter: InMemorySpanExporter) -> None:
"""Test that fan-out edge group processing creates proper success spans."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
@@ -827,7 +830,7 @@ async def test_fan_out_edge_group_tracing_success(span_exporter) -> None:
assert link.context.span_id == int("00f067aa0ba902b7", 16)
-async def test_fan_out_edge_group_tracing_with_target(span_exporter) -> None:
+async def test_fan_out_edge_group_tracing_with_target(span_exporter: InMemorySpanExporter) -> None:
"""Test that fan-out edge group processing creates proper spans for targeted messages."""
source = MockExecutor(id="source_executor")
target1 = MockExecutor(id="target_executor_1")
@@ -994,7 +997,7 @@ async def test_target_edge_group_send_message_with_invalid_data() -> None:
assert success is False
-async def test_fan_in_edge_group_tracing_buffered(span_exporter) -> None:
+async def test_fan_in_edge_group_tracing_buffered(span_exporter: InMemorySpanExporter) -> None:
"""Test that fan-in edge group processing creates proper spans for buffered messages."""
source1 = MockExecutor(id="source_executor_1")
source2 = MockExecutor(id="source_executor_2")
@@ -1086,7 +1089,7 @@ async def test_fan_in_edge_group_tracing_buffered(span_exporter) -> None:
assert link.context.span_id == int("00f067aa0ba902b8", 16)
-async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter) -> None:
+async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter: InMemorySpanExporter) -> None:
"""Test that fan-in edge group processing creates proper spans for type mismatches."""
source1 = MockExecutor(id="source_executor_1")
source2 = MockExecutor(id="source_executor_2")
diff --git a/python/packages/core/tests/workflow/test_executor.py b/python/packages/core/tests/workflow/test_executor.py
index 06d027f19d..77827c0634 100644
--- a/python/packages/core/tests/workflow/test_executor.py
+++ b/python/packages/core/tests/workflow/test_executor.py
@@ -3,8 +3,6 @@
from dataclasses import dataclass
import pytest
-from typing_extensions import Never
-
from agent_framework import (
Executor,
Message,
@@ -16,6 +14,7 @@ from agent_framework import (
handler,
response_handler,
)
+from typing_extensions import Never
# Module-level types for string forward reference tests
@@ -59,7 +58,7 @@ def test_executor_handler_without_annotations():
class MockExecutorWithOneHandlerWithoutAnnotations(Executor): # type: ignore
"""A mock executor with one handler that does not implement any annotations."""
- @handler
+ @handler # pyright: ignore[reportUnknownArgumentType]
async def handle(self, message, ctx) -> None: # type: ignore
"""A mock handler that does not implement any annotations."""
pass
@@ -156,7 +155,11 @@ async def test_executor_invoked_event_contains_input_data():
workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, collector).build()
events = await workflow.run("hello world")
- invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
+ invoked_events = [
+ e
+ for e in events
+ if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
+ ]
assert len(invoked_events) == 2
@@ -190,10 +193,16 @@ async def test_executor_completed_event_contains_sent_messages():
sender = MultiSenderExecutor(id="sender")
collector = CollectorExecutor(id="collector")
- workflow = WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
+ workflow = (
+ WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
+ )
events = await workflow.run("hello")
- completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+ completed_events = [
+ e
+ for e in events
+ if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
+ ]
# Sender should have completed with the sent messages
sender_completed = next(e for e in completed_events if e.executor_id == "sender")
@@ -201,7 +210,9 @@ async def test_executor_completed_event_contains_sent_messages():
assert sender_completed.data == ["hello-first", "hello-second"]
# Collector should have completed with no sent messages (None)
- collector_completed_events = [e for e in completed_events if e.executor_id == "collector"]
+ collector_completed_events = [
+ e for e in completed_events if e.executor_id == "collector"
+ ]
# Collector is called twice (once per message from sender)
assert len(collector_completed_events) == 2
for collector_completed in collector_completed_events:
@@ -220,7 +231,11 @@ async def test_executor_completed_event_includes_yielded_outputs():
workflow = WorkflowBuilder(start_executor=executor).build()
events = await workflow.run("test")
- completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+ completed_events = [
+ e
+ for e in events
+ if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
+ ]
assert len(completed_events) == 1
assert completed_events[0].executor_id == "yielder"
@@ -248,7 +263,9 @@ async def test_executor_events_with_complex_message_types():
class ProcessorExecutor(Executor):
@handler
- async def handle(self, request: Request, ctx: WorkflowContext[Response]) -> None:
+ async def handle(
+ self, request: Request, ctx: WorkflowContext[Response]
+ ) -> None:
response = Response(results=[request.query.upper()] * request.limit)
await ctx.send_message(response)
@@ -260,13 +277,23 @@ async def test_executor_events_with_complex_message_types():
processor = ProcessorExecutor(id="processor")
collector = CollectorExecutor(id="collector")
- workflow = WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
+ workflow = (
+ WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
+ )
input_request = Request(query="hello", limit=3)
events = await workflow.run(input_request)
- invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
- completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
+ invoked_events = [
+ e
+ for e in events
+ if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
+ ]
+ completed_events = [
+ e
+ for e in events
+ if isinstance(e, WorkflowEvent) and e.type == "executor_completed"
+ ]
# Check processor invoked event has the Request object
processor_invoked = next(e for e in invoked_events if e.executor_id == "processor")
@@ -275,7 +302,9 @@ async def test_executor_events_with_complex_message_types():
assert processor_invoked.data.limit == 3
# Check processor completed event has the Response object
- processor_completed = next(e for e in completed_events if e.executor_id == "processor")
+ processor_completed = next(
+ e for e in completed_events if e.executor_id == "processor"
+ )
assert processor_completed.data is not None
assert len(processor_completed.data) == 1
assert isinstance(processor_completed.data[0], Response)
@@ -361,7 +390,9 @@ def test_executor_workflow_output_types_property():
# Test executor with union workflow output types
class UnionWorkflowOutputExecutor(Executor):
@handler
- async def handle(self, text: str, ctx: WorkflowContext[int, str | bool]) -> None:
+ async def handle(
+ self, text: str, ctx: WorkflowContext[int, str | bool]
+ ) -> None:
pass
executor = UnionWorkflowOutputExecutor(id="union_workflow_output")
@@ -372,11 +403,15 @@ def test_executor_workflow_output_types_property():
# Test executor with multiple handlers having different workflow output types
class MultiHandlerWorkflowExecutor(Executor):
@handler
- async def handle_string(self, text: str, ctx: WorkflowContext[int, str]) -> None:
+ async def handle_string(
+ self, text: str, ctx: WorkflowContext[int, str]
+ ) -> None:
pass
@handler
- async def handle_number(self, num: int, ctx: WorkflowContext[bool, float]) -> None:
+ async def handle_number(
+ self, num: int, ctx: WorkflowContext[bool, float]
+ ) -> None:
pass
executor = MultiHandlerWorkflowExecutor(id="multi_workflow")
@@ -430,7 +465,9 @@ def test_executor_output_types_includes_response_handlers():
pass
@response_handler
- async def handle_response(self, original_request: str, response: bool, ctx: WorkflowContext[float]) -> None:
+ async def handle_response(
+ self, original_request: str, response: bool, ctx: WorkflowContext[float]
+ ) -> None:
pass
executor = RequestResponseExecutor(id="request_response")
@@ -452,7 +489,10 @@ def test_executor_workflow_output_types_includes_response_handlers():
@response_handler
async def handle_response(
- self, original_request: str, response: bool, ctx: WorkflowContext[float, bool]
+ self,
+ original_request: str,
+ response: bool,
+ ctx: WorkflowContext[float, bool],
) -> None:
pass
@@ -509,7 +549,10 @@ def test_executor_response_handler_union_output_types():
@response_handler
async def handle_response(
- self, original_request: str, response: bool, ctx: WorkflowContext[int | str | float, bool | int]
+ self,
+ original_request: str,
+ response: bool,
+ ctx: WorkflowContext[int | str | float, bool | int],
) -> None:
pass
@@ -531,7 +574,9 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
"""Test that executor_invoked event (type='executor_invoked').data captures original input, not mutated input."""
@executor(id="Mutator")
- async def mutator(messages: list[Message], ctx: WorkflowContext[list[Message]]) -> None:
+ async def mutator(
+ messages: list[Message], ctx: WorkflowContext[list[Message]]
+ ) -> None:
# The handler mutates the input list by appending new messages
original_len = len(messages)
messages.append(Message(role="assistant", text="Added by executor"))
@@ -546,7 +591,11 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
events = await workflow.run(input_messages)
# Find the invoked event for the Mutator executor
- invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
+ invoked_events = [
+ e
+ for e in events
+ if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"
+ ]
assert len(invoked_events) == 1
mutator_invoked = invoked_events[0]
@@ -577,8 +626,8 @@ class TestHandlerExplicitTypes:
exec_instance = ExplicitInputExecutor(id="explicit_input")
# Handler should be registered for str (explicit), not Any (introspected)
- assert str in exec_instance._handlers
- assert len(exec_instance._handlers) == 1
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ assert len(exec_instance._handlers) == 1 # pyright: ignore[reportPrivateUsage]
# Can handle str messages
assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
@@ -596,8 +645,8 @@ class TestHandlerExplicitTypes:
exec_instance = ExplicitOutputExecutor(id="explicit_output")
# Handler spec should have int as output type (explicit)
- handler_func = exec_instance._handlers[str]
- assert handler_func._handler_spec["output_types"] == [int]
+ handler_func = exec_instance._handlers[str] # pyright: ignore[reportPrivateUsage]
+ assert handler_func._handler_spec["output_types"] == [int] # pyright: ignore[reportFunctionMemberAccess]
# Executor output_types property should reflect explicit type
assert int in exec_instance.output_types
@@ -615,16 +664,20 @@ class TestHandlerExplicitTypes:
exec_instance = ExplicitBothExecutor(id="explicit_both")
# Handler should be registered for dict (explicit input type)
- assert dict in exec_instance._handlers
- assert len(exec_instance._handlers) == 1
+ assert dict in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ assert len(exec_instance._handlers) == 1 # pyright: ignore[reportPrivateUsage]
# Output type should be list (explicit)
- handler_func = exec_instance._handlers[dict]
- assert handler_func._handler_spec["output_types"] == [list]
+ handler_func = exec_instance._handlers[dict] # pyright: ignore[reportPrivateUsage]
+ assert handler_func._handler_spec["output_types"] == [list] # pyright: ignore[reportFunctionMemberAccess]
# Verify can_handle
- assert exec_instance.can_handle(WorkflowMessage(data={"key": "value"}, source_id="mock"))
- assert not exec_instance.can_handle(WorkflowMessage(data="string", source_id="mock"))
+ assert exec_instance.can_handle(
+ WorkflowMessage(data={"key": "value"}, source_id="mock")
+ )
+ assert not exec_instance.can_handle(
+ WorkflowMessage(data="string", source_id="mock")
+ )
def test_handler_with_explicit_union_input_type(self):
"""Test that explicit union input_type is handled correctly."""
@@ -639,13 +692,15 @@ class TestHandlerExplicitTypes:
# Handler should be registered for the union type
# The union type itself is stored as the key
- assert len(exec_instance._handlers) == 1
+ assert len(exec_instance._handlers) == 1 # pyright: ignore[reportPrivateUsage]
# Can handle both str and int messages
assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
assert exec_instance.can_handle(WorkflowMessage(data=42, source_id="mock"))
# Cannot handle float
- assert not exec_instance.can_handle(WorkflowMessage(data=3.14, source_id="mock"))
+ assert not exec_instance.can_handle(
+ WorkflowMessage(data=3.14, source_id="mock")
+ )
def test_handler_with_explicit_union_output_type(self):
"""Test that explicit union output is normalized to a list."""
@@ -674,8 +729,8 @@ class TestHandlerExplicitTypes:
exec_instance = PrecedenceExecutor(id="precedence")
# Should use explicit input type (bytes), not introspected (str)
- assert bytes in exec_instance._handlers
- assert str not in exec_instance._handlers
+ assert bytes in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ assert str not in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
# Should use explicit output type (float), not introspected (int)
assert float in exec_instance.output_types
@@ -692,7 +747,7 @@ class TestHandlerExplicitTypes:
exec_instance = IntrospectedExecutor(id="introspected")
# Should use introspected types
- assert str in exec_instance._handlers
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
assert int in exec_instance.output_types
def test_handler_explicit_mode_requires_input(self):
@@ -705,13 +760,13 @@ class TestHandlerExplicitTypes:
pass
exec_input = OnlyInputExecutor(id="only_input")
- assert bytes in exec_input._handlers # Explicit
+ assert bytes in exec_input._handlers # pyright: ignore[reportPrivateUsage] # Explicit
assert exec_input.output_types == [] # No output types (not introspected)
# Only explicit output without input should raise error
with pytest.raises(ValueError, match="must specify 'input' type"):
- class OnlyOutputExecutor(Executor):
+ class OnlyOutputExecutor(Executor): # pyright: ignore[reportUnusedClass]
@handler(output=float)
async def handle(self, message: str, ctx: WorkflowContext[int]) -> None:
pass
@@ -719,9 +774,11 @@ class TestHandlerExplicitTypes:
# Only explicit workflow_output without input should raise error
with pytest.raises(ValueError, match="must specify 'input' type"):
- class OnlyWorkflowOutputExecutor(Executor):
+ class OnlyWorkflowOutputExecutor(Executor): # pyright: ignore[reportUnusedClass]
@handler(workflow_output=bool)
- async def handle(self, message: str, ctx: WorkflowContext[int, str]) -> None:
+ async def handle(
+ self, message: str, ctx: WorkflowContext[int, str]
+ ) -> None:
pass
def test_handler_explicit_input_type_allows_no_message_annotation(self):
@@ -734,8 +791,7 @@ class TestHandlerExplicitTypes:
exec_instance = NoAnnotationExecutor(id="no_annotation")
- # Should work with explicit input_type
- assert str in exec_instance._handlers
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
def test_handler_multiple_handlers_mixed_explicit_and_introspected(self):
@@ -747,15 +803,17 @@ class TestHandlerExplicitTypes:
pass
@handler
- async def handle_introspected(self, message: float, ctx: WorkflowContext[bool]) -> None:
+ async def handle_introspected(
+ self, message: float, ctx: WorkflowContext[bool]
+ ) -> None:
pass
exec_instance = MixedExecutor(id="mixed")
# Should have both handlers
- assert len(exec_instance._handlers) == 2
- assert str in exec_instance._handlers # Explicit
- assert float in exec_instance._handlers # Introspected
+ assert len(exec_instance._handlers) == 2 # pyright: ignore[reportPrivateUsage]
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage] # Explicit
+ assert float in exec_instance._handlers # pyright: ignore[reportPrivateUsage] # Introspected
# Should have both output types
assert int in exec_instance.output_types # Explicit
@@ -772,8 +830,10 @@ class TestHandlerExplicitTypes:
exec_instance = StringRefExecutor(id="string_ref")
# Should resolve the string to the actual type
- assert ForwardRefMessage in exec_instance._handlers
- assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock"))
+ assert ForwardRefMessage in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ assert exec_instance.can_handle(
+ WorkflowMessage(data=ForwardRefMessage("hello"), source_id="mock")
+ )
def test_handler_with_string_forward_reference_union(self):
"""Test that string forward references work with union types."""
@@ -786,8 +846,12 @@ class TestHandlerExplicitTypes:
exec_instance = StringUnionExecutor(id="string_union")
# Should handle both types
- assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock"))
- assert exec_instance.can_handle(WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock"))
+ assert exec_instance.can_handle(
+ WorkflowMessage(data=ForwardRefTypeA("hello"), source_id="mock")
+ )
+ assert exec_instance.can_handle(
+ WorkflowMessage(data=ForwardRefTypeB(42), source_id="mock")
+ )
def test_handler_with_string_forward_reference_output_type(self):
"""Test that string forward references work for output_type."""
@@ -813,8 +877,8 @@ class TestHandlerExplicitTypes:
exec_instance = ExplicitWorkflowOutputExecutor(id="explicit_workflow_output")
# Handler spec should have bool as workflow_output_type (explicit)
- handler_func = exec_instance._handlers[str]
- assert handler_func._handler_spec["workflow_output_types"] == [bool]
+ handler_func = exec_instance._handlers[str] # pyright: ignore[reportPrivateUsage]
+ assert handler_func._handler_spec["workflow_output_types"] == [bool] # pyright: ignore[reportFunctionMemberAccess]
# Executor workflow_output_types property should reflect explicit type
assert bool in exec_instance.workflow_output_types
@@ -826,13 +890,14 @@ class TestHandlerExplicitTypes:
class PrecedenceExecutor(Executor):
@handler(input=int, output=float, workflow_output=str)
- async def handle(self, message: int, ctx: WorkflowContext[int, bool]) -> None:
+ async def handle(
+ self, message: int, ctx: WorkflowContext[int, bool]
+ ) -> None:
pass
exec_instance = PrecedenceExecutor(id="precedence")
- # All types should come from explicit params
- assert int in exec_instance._handlers
+ assert int in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
assert float in exec_instance.output_types
assert str in exec_instance.workflow_output_types
# Introspected types should NOT be present
@@ -849,8 +914,7 @@ class TestHandlerExplicitTypes:
exec_instance = AllExplicitExecutor(id="all_explicit")
- # Check input type
- assert str in exec_instance._handlers
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
assert exec_instance.can_handle(WorkflowMessage(data="hello", source_id="mock"))
# Check output_type
@@ -894,7 +958,9 @@ class TestHandlerExplicitTypes:
async def handle(self, message, ctx: WorkflowContext) -> None: # type: ignore[no-untyped-def]
pass
- exec_instance = StringUnionWorkflowOutputExecutor(id="string_union_workflow_output")
+ exec_instance = StringUnionWorkflowOutputExecutor(
+ id="string_union_workflow_output"
+ )
# Should resolve both types from string union
assert ForwardRefTypeA in exec_instance.workflow_output_types
@@ -905,10 +971,14 @@ class TestHandlerExplicitTypes:
class IntrospectedWorkflowOutputExecutor(Executor):
@handler
- async def handle(self, message: str, ctx: WorkflowContext[int, bool]) -> None:
+ async def handle(
+ self, message: str, ctx: WorkflowContext[int, bool]
+ ) -> None:
pass
- exec_instance = IntrospectedWorkflowOutputExecutor(id="introspected_workflow_output")
+ exec_instance = IntrospectedWorkflowOutputExecutor(
+ id="introspected_workflow_output"
+ )
# Should use introspected types from WorkflowContext[int, bool]
assert int in exec_instance.output_types
diff --git a/python/packages/core/tests/workflow/test_executor_future.py b/python/packages/core/tests/workflow/test_executor_future.py
index c0916b9cf7..cb0c5c9f58 100644
--- a/python/packages/core/tests/workflow/test_executor_future.py
+++ b/python/packages/core/tests/workflow/test_executor_future.py
@@ -34,8 +34,8 @@ class TestExecutorFutureAnnotations:
pass
exec_instance = MyExecutor(id="test")
- assert str in exec_instance._handlers
- spec = exec_instance._handler_specs[0]
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ spec = exec_instance._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is str
assert spec["output_types"] == [MyTypeA]
assert spec["workflow_output_types"] == [MyTypeB]
@@ -49,8 +49,8 @@ class TestExecutorFutureAnnotations:
pass
exec_instance = MyExecutor(id="test")
- assert int in exec_instance._handlers
- spec = exec_instance._handler_specs[0]
+ assert int in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ spec = exec_instance._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is int
assert spec["output_types"] == [MyTypeA]
@@ -63,7 +63,7 @@ class TestExecutorFutureAnnotations:
pass
exec_instance = MyExecutor(id="test")
- spec = exec_instance._handler_specs[0]
+ spec = exec_instance._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] == dict[str, Any]
assert spec["output_types"] == [list[str]]
@@ -76,8 +76,8 @@ class TestExecutorFutureAnnotations:
pass
exec_instance = MyExecutor(id="test")
- assert str in exec_instance._handlers
- spec = exec_instance._handler_specs[0]
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ spec = exec_instance._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["output_types"] == []
assert spec["workflow_output_types"] == []
@@ -86,12 +86,12 @@ class TestExecutorFutureAnnotations:
class MyExecutor(Executor):
@handler(input=str, output=MyTypeA)
- async def example(self, input, ctx) -> None:
+ async def example(self, input, ctx) -> None: # type: ignore[no-untyped-def]
pass
exec_instance = MyExecutor(id="test")
- assert str in exec_instance._handlers
- spec = exec_instance._handler_specs[0]
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ spec = exec_instance._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is str
assert spec["output_types"] == [MyTypeA]
@@ -104,8 +104,8 @@ class TestExecutorFutureAnnotations:
pass
exec_instance = MyExecutor(id="test")
- assert str in exec_instance._handlers
- spec = exec_instance._handler_specs[0]
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
+ spec = exec_instance._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["output_types"] == [MyTypeA, MyTypeB]
assert spec["workflow_output_types"] == [MyTypeC]
@@ -118,7 +118,7 @@ class TestExecutorFutureAnnotations:
"""
with pytest.raises(ValueError):
- class Bad(Executor):
- @handler
- async def example(self, input: NonExistentType, ctx: WorkflowContext[MyTypeA, MyTypeB]) -> None: # noqa: F821
+ class Bad(Executor): # pyright: ignore[reportUnusedClass]
+ @handler # pyright: ignore[reportUnknownArgumentType]
+ async def example(self, input: NonExistentType, ctx: WorkflowContext[MyTypeA, MyTypeB]) -> None: # noqa: F821 # type: ignore[name-defined]
pass
diff --git a/python/packages/core/tests/workflow/test_full_conversation.py b/python/packages/core/tests/workflow/test_full_conversation.py
index 20d9abd8c0..b6b5260d83 100644
--- a/python/packages/core/tests/workflow/test_full_conversation.py
+++ b/python/packages/core/tests/workflow/test_full_conversation.py
@@ -1,7 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
-from collections.abc import AsyncIterable, Awaitable, Sequence
-from typing import Any
+from collections.abc import AsyncIterable, Awaitable
+from typing import Any, Literal, overload
import pytest
from pydantic import PrivateAttr
@@ -13,6 +13,7 @@ from agent_framework import (
AgentExecutorResponse,
AgentResponse,
AgentResponseUpdate,
+ AgentRunInputs,
AgentSession,
BaseAgent,
Content,
@@ -34,14 +35,32 @@ class _SimpleAgent(BaseAgent):
super().__init__(**kwargs)
self._reply_text = reply_text
+ @overload
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -81,14 +100,32 @@ class _ToolHistoryAgent(BaseAgent):
Message(role="assistant", contents=[Content.from_text(text=self._summary_text)]),
]
+ @overload
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
@@ -165,14 +202,32 @@ class _CaptureAgent(BaseAgent):
super().__init__(**kwargs)
self._reply_text = reply_text
+ @overload
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
# Normalize and record messages for verification
norm: list[Message] = []
if messages:
@@ -260,7 +315,7 @@ class _RoundTripCoordinator(Executor):
async def handle_response(
self,
response: AgentExecutorResponse,
- ctx: WorkflowContext[Never, dict[str, Any]],
+ ctx: WorkflowContext[AgentExecutorRequest, dict[str, Any]],
) -> None:
self._seen += 1
if self._seen == 1:
@@ -314,14 +369,32 @@ class _SessionIdCapturingAgent(BaseAgent):
_captured_service_session_id: str | None = PrivateAttr(default="NOT_CAPTURED")
+ @overload
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self._captured_service_session_id = session.service_session_id if session else None
async def _run() -> AgentResponse:
@@ -342,7 +415,7 @@ class _FullHistoryReplayCoordinator(Executor):
async def handle(
self,
response: AgentExecutorResponse,
- ctx: WorkflowContext[Never, Any],
+ ctx: WorkflowContext[AgentExecutorRequest, Any],
) -> None:
full_conv = list(response.full_conversation or response.agent_response.messages)
full_conv.append(Message(role="user", text="follow-up"))
diff --git a/python/packages/core/tests/workflow/test_function_executor.py b/python/packages/core/tests/workflow/test_function_executor.py
index c0b73156ff..8bb3f94d29 100644
--- a/python/packages/core/tests/workflow/test_function_executor.py
+++ b/python/packages/core/tests/workflow/test_function_executor.py
@@ -48,12 +48,12 @@ class TestFunctionExecutor:
func_exec = FunctionExecutor(process_string)
# Check that handler was registered
- assert len(func_exec._handlers) == 1
- assert str in func_exec._handlers
+ assert len(func_exec._handlers) == 1 # pyright: ignore[reportPrivateUsage]
+ assert str in func_exec._handlers # pyright: ignore[reportPrivateUsage]
# Check handler spec was created
- assert len(func_exec._handler_specs) == 1
- spec = func_exec._handler_specs[0]
+ assert len(func_exec._handler_specs) == 1 # pyright: ignore[reportPrivateUsage]
+ spec = func_exec._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["name"] == "process_string"
assert spec["message_type"] is str
assert spec["output_types"] == [str]
@@ -67,10 +67,10 @@ class TestFunctionExecutor:
assert isinstance(process_int, FunctionExecutor)
assert process_int.id == "test_executor"
- assert int in process_int._handlers
+ assert int in process_int._handlers # pyright: ignore[reportPrivateUsage]
# Check spec
- spec = process_int._handler_specs[0]
+ spec = process_int._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is int
assert spec["output_types"] == [int]
@@ -78,7 +78,7 @@ class TestFunctionExecutor:
"""Test @executor decorator uses function name as default ID."""
@executor
- async def my_function(data: dict, ctx: WorkflowContext[Any]) -> None:
+ async def my_function(data: dict[str, Any], ctx: WorkflowContext[Any]) -> None:
await ctx.send_message(data)
assert my_function.id == "my_function"
@@ -92,7 +92,7 @@ class TestFunctionExecutor:
assert isinstance(no_parens_function, FunctionExecutor)
assert no_parens_function.id == "no_parens_function"
- assert str in no_parens_function._handlers
+ assert str in no_parens_function._handlers # pyright: ignore[reportPrivateUsage]
# Also test with single parameter function
@executor
@@ -101,7 +101,7 @@ class TestFunctionExecutor:
assert isinstance(simple_no_parens, FunctionExecutor)
assert simple_no_parens.id == "simple_no_parens"
- assert int in simple_no_parens._handlers
+ assert int in simple_no_parens._handlers # pyright: ignore[reportPrivateUsage]
def test_union_output_types(self):
"""Test that union output types are properly inferred for both messages and workflow outputs."""
@@ -113,7 +113,7 @@ class TestFunctionExecutor:
else:
await ctx.send_message(text.upper())
- spec = multi_output._handler_specs[0]
+ spec = multi_output._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert set(spec["output_types"]) == {str, int}
assert spec["workflow_output_types"] == [] # No workflow outputs defined
@@ -127,7 +127,7 @@ class TestFunctionExecutor:
else:
await ctx.yield_output(data.upper())
- workflow_spec = multi_workflow_output._handler_specs[0]
+ workflow_spec = multi_workflow_output._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert workflow_spec["output_types"] == [] # None means no message outputs
assert set(workflow_spec["workflow_output_types"]) == {str, int, bool}
@@ -139,7 +139,7 @@ class TestFunctionExecutor:
# This executor doesn't send any messages
pass
- spec = no_output._handler_specs[0]
+ spec = no_output._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["output_types"] == []
assert spec["workflow_output_types"] == [] # No workflow outputs defined
@@ -150,7 +150,7 @@ class TestFunctionExecutor:
async def any_output(data: str, ctx: WorkflowContext[Any]) -> None:
await ctx.send_message("result")
- spec = any_output._handler_specs[0]
+ spec = any_output._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["output_types"] == [Any]
assert spec["workflow_output_types"] == [] # No workflow outputs defined
@@ -160,7 +160,7 @@ class TestFunctionExecutor:
await ctx.send_message("message")
await ctx.yield_output("workflow_output")
- both_spec = any_both_output._handler_specs[0]
+ both_spec = any_both_output._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert both_spec["output_types"] == [Any]
assert both_spec["workflow_output_types"] == [Any]
@@ -228,11 +228,11 @@ class TestFunctionExecutor:
await ctx.yield_output(result)
# Verify type inference for both executors
- upper_spec = to_upper._handler_specs[0]
+ upper_spec = to_upper._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert upper_spec["output_types"] == [str]
assert upper_spec["workflow_output_types"] == [] # No workflow outputs
- reverse_spec = reverse_text._handler_specs[0]
+ reverse_spec = reverse_text._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert reverse_spec["output_types"] == [Any] # First parameter is Any
assert reverse_spec["workflow_output_types"] == [str] # Second parameter is str
@@ -270,7 +270,7 @@ class TestFunctionExecutor:
await ctx.send_message(message)
with pytest.raises(ValueError, match="Handler for type .* already registered"):
- func_exec._register_instance_handler(
+ func_exec._register_instance_handler( # pyright: ignore[reportPrivateUsage]
name="second",
func=second_handler,
message_type=str,
@@ -287,7 +287,7 @@ class TestFunctionExecutor:
result = {item: len(item) for item in items}
await ctx.send_message(result)
- spec = process_list._handler_specs[0]
+ spec = process_list._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] == list[str]
assert spec["output_types"] == [dict[str, int]]
@@ -300,10 +300,10 @@ class TestFunctionExecutor:
assert isinstance(process_simple, FunctionExecutor)
assert process_simple.id == "simple_processor"
- assert str in process_simple._handlers
+ assert str in process_simple._handlers # pyright: ignore[reportPrivateUsage]
# Check spec - single parameter functions have no output types since they can't send messages
- spec = process_simple._handler_specs[0]
+ spec = process_simple._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is str
assert spec["output_types"] == []
assert spec["ctx_annotation"] is None
@@ -316,7 +316,7 @@ class TestFunctionExecutor:
return data * 2
func_exec = FunctionExecutor(valid_single)
- assert int in func_exec._handlers
+ assert int in func_exec._handlers # pyright: ignore[reportPrivateUsage]
# Single parameter with missing type annotation should still fail
async def no_annotation(data): # type: ignore
@@ -349,7 +349,7 @@ class TestFunctionExecutor:
# For testing purposes, we can check that the handler is registered correctly
assert double_value.can_handle(WorkflowMessage(data=5, source_id="mock"))
- assert int in double_value._handlers
+ assert int in double_value._handlers # pyright: ignore[reportPrivateUsage]
def test_sync_function_basic(self):
"""Test basic synchronous function support."""
@@ -360,10 +360,10 @@ class TestFunctionExecutor:
assert isinstance(process_sync, FunctionExecutor)
assert process_sync.id == "sync_processor"
- assert str in process_sync._handlers
+ assert str in process_sync._handlers # pyright: ignore[reportPrivateUsage]
# Check spec - sync single parameter functions have no output types
- spec = process_sync._handler_specs[0]
+ spec = process_sync._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is str
assert spec["output_types"] == []
assert spec["ctx_annotation"] is None
@@ -378,10 +378,10 @@ class TestFunctionExecutor:
assert isinstance(sync_with_ctx, FunctionExecutor)
assert sync_with_ctx.id == "sync_with_ctx"
- assert int in sync_with_ctx._handlers
+ assert int in sync_with_ctx._handlers # pyright: ignore[reportPrivateUsage]
# Check spec - sync functions with context can infer output types
- spec = sync_with_ctx._handler_specs[0]
+ spec = sync_with_ctx._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is int
assert spec["output_types"] == [int]
@@ -404,18 +404,18 @@ class TestFunctionExecutor:
return data.upper()
func_exec = FunctionExecutor(valid_sync)
- assert str in func_exec._handlers
+ assert str in func_exec._handlers # pyright: ignore[reportPrivateUsage]
# Valid sync function with two parameters
def valid_sync_with_ctx(data: int, ctx: WorkflowContext[str]):
return str(data)
func_exec2 = FunctionExecutor(valid_sync_with_ctx)
- assert int in func_exec2._handlers
+ assert int in func_exec2._handlers # pyright: ignore[reportPrivateUsage]
# Sync function with missing type annotation should still fail
- def no_annotation(data): # type: ignore
- return data
+ def no_annotation(data): # type: ignore # pyright: ignore[reportUnknownVariableType]
+ return data # pyright: ignore[reportUnknownVariableType]
with pytest.raises(ValueError, match="type annotation for the message"):
FunctionExecutor(no_annotation) # type: ignore
@@ -457,11 +457,11 @@ class TestFunctionExecutor:
await ctx.yield_output(result)
# Verify type inference for sync and async functions
- sync_spec = to_upper_sync._handler_specs[0]
+ sync_spec = to_upper_sync._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert sync_spec["output_types"] == [str]
assert sync_spec["workflow_output_types"] == [] # No workflow outputs
- async_spec = reverse_async._handler_specs[0]
+ async_spec = reverse_async._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert async_spec["output_types"] == [Any] # First parameter is Any
assert async_spec["workflow_output_types"] == [str] # Second parameter is str
@@ -471,8 +471,8 @@ class TestFunctionExecutor:
# For integration testing, we mainly verify that the handlers are properly registered
# and the functions are wrapped correctly
- assert str in to_upper_sync._handlers
- assert str in reverse_async._handlers
+ assert str in to_upper_sync._handlers # pyright: ignore[reportPrivateUsage]
+ assert str in reverse_async._handlers # pyright: ignore[reportPrivateUsage]
async def test_sync_function_thread_execution(self):
"""Test that sync functions run in thread pool and don't block the event loop."""
@@ -491,13 +491,13 @@ class TestFunctionExecutor:
return data.upper()
# Verify the function is wrapped and registered
- assert str in blocking_function._handlers
+ assert str in blocking_function._handlers # pyright: ignore[reportPrivateUsage]
# For a more complete test, we'd need to create a full workflow context,
# but for now we can verify that the function was properly wrapped
# and that sync functions store the correct metadata
- assert not blocking_function._is_async
- assert not blocking_function._has_context
+ assert not blocking_function._is_async # pyright: ignore[reportPrivateUsage]
+ assert not blocking_function._has_context # pyright: ignore[reportPrivateUsage]
# The actual thread execution test would require a full workflow setup,
# but the important thing is that asyncio.to_thread is used in the wrapper
@@ -506,7 +506,7 @@ class TestFunctionExecutor:
"""Test that @executor decorator properly rejects @staticmethod with clear error."""
with pytest.raises(ValueError) as exc_info:
- class Example:
+ class Example: # pyright: ignore[reportUnusedClass]
@executor
@staticmethod
async def bad_handler(data: str) -> str:
@@ -519,7 +519,7 @@ class TestFunctionExecutor:
"""Test that @executor decorator properly rejects @classmethod with clear error."""
with pytest.raises(ValueError) as exc_info:
- class Example:
+ class Example: # pyright: ignore[reportUnusedClass]
@executor
@classmethod
async def bad_handler(cls, data: str) -> str:
@@ -570,8 +570,8 @@ class TestExecutorExplicitTypes:
pass
# Handler should be registered for str (explicit)
- assert str in process._handlers
- assert len(process._handlers) == 1
+ assert str in process._handlers # pyright: ignore[reportPrivateUsage]
+ assert len(process._handlers) == 1 # pyright: ignore[reportPrivateUsage]
# Can handle str messages
assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
@@ -586,7 +586,7 @@ class TestExecutorExplicitTypes:
pass
# Handler spec should have int as output type (explicit), not str (introspected)
- spec = process._handler_specs[0]
+ spec = process._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["output_types"] == [int]
# Executor output_types property should reflect explicit type
@@ -601,11 +601,11 @@ class TestExecutorExplicitTypes:
pass
# Handler should be registered for dict (explicit input type)
- assert dict in process._handlers
- assert len(process._handlers) == 1
+ assert dict in process._handlers # pyright: ignore[reportPrivateUsage]
+ assert len(process._handlers) == 1 # pyright: ignore[reportPrivateUsage]
# Output type should be list (explicit)
- spec = process._handler_specs[0]
+ spec = process._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["output_types"] == [list]
# Verify can_handle
@@ -620,7 +620,7 @@ class TestExecutorExplicitTypes:
pass
# Handler should be registered for the union type
- assert len(process._handlers) == 1
+ assert len(process._handlers) == 1 # pyright: ignore[reportPrivateUsage]
# Can handle both str and int messages
assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
@@ -648,8 +648,8 @@ class TestExecutorExplicitTypes:
pass
# Should use explicit input type (bytes), not introspected (str)
- assert bytes in process._handlers
- assert str not in process._handlers
+ assert bytes in process._handlers # pyright: ignore[reportPrivateUsage]
+ assert str not in process._handlers # pyright: ignore[reportPrivateUsage]
# Should use explicit output type (float), not introspected (int)
assert float in process.output_types
@@ -663,7 +663,7 @@ class TestExecutorExplicitTypes:
pass
# Should use introspected types
- assert str in process._handlers
+ assert str in process._handlers # pyright: ignore[reportPrivateUsage]
assert int in process.output_types
def test_executor_partial_explicit_types(self):
@@ -674,7 +674,7 @@ class TestExecutorExplicitTypes:
async def process_input(message: str, ctx: WorkflowContext[int]) -> None:
pass
- assert bytes in process_input._handlers # Explicit
+ assert bytes in process_input._handlers # Explicit # pyright: ignore[reportPrivateUsage]
assert int in process_input.output_types # Introspected
# Only explicit output_type, introspect input_type
@@ -682,7 +682,7 @@ class TestExecutorExplicitTypes:
async def process_output(message: str, ctx: WorkflowContext[int]) -> None:
pass
- assert str in process_output._handlers # Introspected
+ assert str in process_output._handlers # Introspected # pyright: ignore[reportPrivateUsage]
assert float in process_output.output_types # Explicit
assert int not in process_output.output_types # Not introspected when explicit provided
@@ -694,7 +694,7 @@ class TestExecutorExplicitTypes:
pass
# Should work with explicit input_type
- assert str in process._handlers
+ assert str in process._handlers # pyright: ignore[reportPrivateUsage]
assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
def test_executor_explicit_types_with_id(self):
@@ -705,7 +705,7 @@ class TestExecutorExplicitTypes:
pass
assert process.id == "custom_id"
- assert bytes in process._handlers
+ assert bytes in process._handlers # pyright: ignore[reportPrivateUsage]
assert int in process.output_types
def test_executor_explicit_types_with_single_param_function(self):
@@ -713,10 +713,10 @@ class TestExecutorExplicitTypes:
@executor(input=str)
async def process(message): # type: ignore[no-untyped-def]
- return message.upper()
+ return message.upper() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType]
# Should work with explicit input_type
- assert str in process._handlers
+ assert str in process._handlers # pyright: ignore[reportPrivateUsage]
assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
assert not process.can_handle(WorkflowMessage(data=42, source_id="mock"))
@@ -727,7 +727,7 @@ class TestExecutorExplicitTypes:
def process(message, ctx: WorkflowContext) -> None: # type: ignore[no-untyped-def]
pass
- assert int in process._handlers
+ assert int in process._handlers # pyright: ignore[reportPrivateUsage]
assert str in process.output_types
def test_function_executor_constructor_with_explicit_types(self):
@@ -736,10 +736,10 @@ class TestExecutorExplicitTypes:
async def process(message, ctx: WorkflowContext) -> None: # type: ignore[no-untyped-def]
pass
- func_exec = FunctionExecutor(process, id="test", input=dict, output=list)
+ func_exec = FunctionExecutor(process, id="test", input=dict, output=list) # pyright: ignore[reportUnknownArgumentType]
- assert dict in func_exec._handlers
- spec = func_exec._handler_specs[0]
+ assert dict in func_exec._handlers # pyright: ignore[reportPrivateUsage]
+ spec = func_exec._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is dict
assert spec["output_types"] == [list]
@@ -766,7 +766,7 @@ class TestExecutorExplicitTypes:
pass
# Should resolve the string to the actual type
- assert FuncExecForwardRefMessage in process._handlers
+ assert FuncExecForwardRefMessage in process._handlers # pyright: ignore[reportPrivateUsage]
assert process.can_handle(WorkflowMessage(data=FuncExecForwardRefMessage("hello"), source_id="mock"))
def test_executor_with_string_forward_reference_union(self):
@@ -798,7 +798,7 @@ class TestExecutorExplicitTypes:
pass
# Handler spec should have bool as workflow_output_type (explicit)
- spec = process._handler_specs[0]
+ spec = process._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["workflow_output_types"] == [bool]
# Executor workflow_output_types property should reflect explicit type
@@ -826,7 +826,7 @@ class TestExecutorExplicitTypes:
pass
# Check input type
- assert str in process._handlers
+ assert str in process._handlers # pyright: ignore[reportPrivateUsage]
assert process.can_handle(WorkflowMessage(data="hello", source_id="mock"))
# Check output_type
@@ -892,6 +892,6 @@ class TestExecutorExplicitTypes:
workflow_output=bool,
)
- assert str in exec_instance._handlers
+ assert str in exec_instance._handlers # pyright: ignore[reportPrivateUsage]
assert int in exec_instance.output_types
assert bool in exec_instance.workflow_output_types
diff --git a/python/packages/core/tests/workflow/test_function_executor_future.py b/python/packages/core/tests/workflow/test_function_executor_future.py
index a4a15aeba0..6d1ed32348 100644
--- a/python/packages/core/tests/workflow/test_function_executor_future.py
+++ b/python/packages/core/tests/workflow/test_function_executor_future.py
@@ -19,10 +19,10 @@ class TestFunctionExecutorFutureAnnotations:
assert isinstance(process_future, FunctionExecutor)
assert process_future.id == "future_test"
- assert int in process_future._handlers
+ assert int in process_future._handlers # pyright: ignore[reportPrivateUsage]
# Check spec
- spec = process_future._handler_specs[0]
+ spec = process_future._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] is int
assert spec["output_types"] == [int]
@@ -34,6 +34,6 @@ class TestFunctionExecutorFutureAnnotations:
await ctx.send_message(["done"])
assert isinstance(process_complex, FunctionExecutor)
- spec = process_complex._handler_specs[0]
+ spec = process_complex._handler_specs[0] # pyright: ignore[reportPrivateUsage]
assert spec["message_type"] == dict[str, Any]
assert spec["output_types"] == [list[str]]
diff --git a/python/packages/core/tests/workflow/test_request_info_mixin.py b/python/packages/core/tests/workflow/test_request_info_mixin.py
index 4c3d6560aa..cfde71b481 100644
--- a/python/packages/core/tests/workflow/test_request_info_mixin.py
+++ b/python/packages/core/tests/workflow/test_request_info_mixin.py
@@ -794,7 +794,7 @@ class TestResponseHandlerExplicitTypes:
"""Test response_handler with explicit request and response types."""
@response_handler(request=str, response=int)
- async def test_handler(self, original_request, response, ctx) -> None:
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
pass
spec = test_handler._response_handler_spec # type: ignore[reportAttributeAccessIssue]
@@ -806,7 +806,7 @@ class TestResponseHandlerExplicitTypes:
"""Test response_handler with explicit output and workflow_output types."""
@response_handler(request=str, response=int, output=bool, workflow_output=float)
- async def test_handler(self, original_request, response, ctx) -> None:
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
pass
spec = test_handler._response_handler_spec # type: ignore[reportAttributeAccessIssue]
@@ -818,8 +818,8 @@ class TestResponseHandlerExplicitTypes:
def test_response_handler_with_union_types(self):
"""Test response_handler with union types."""
- @response_handler(request=str | int, response=bool | float)
- async def test_handler(self, original_request, response, ctx) -> None:
+ @response_handler(request=str | int, response=bool | float) # pyright: ignore[reportArgumentType]
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
pass
spec = test_handler._response_handler_spec # type: ignore[reportAttributeAccessIssue]
@@ -830,7 +830,7 @@ class TestResponseHandlerExplicitTypes:
"""Test response_handler with string forward references."""
@response_handler(request="str", response="int")
- async def test_handler(self, original_request, response, ctx) -> None:
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
pass
spec = test_handler._response_handler_spec # type: ignore[reportAttributeAccessIssue]
@@ -842,7 +842,7 @@ class TestResponseHandlerExplicitTypes:
with pytest.raises(ValueError, match="must specify 'request' type"):
@response_handler(response=int)
- async def test_handler(self, original_request, response, ctx) -> None:
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None: # pyright: ignore[reportUnusedFunction]
pass
def test_response_handler_explicit_missing_response_raises_error(self):
@@ -850,7 +850,7 @@ class TestResponseHandlerExplicitTypes:
with pytest.raises(ValueError, match="must specify 'response' type"):
@response_handler(request=str)
- async def test_handler(self, original_request, response, ctx) -> None:
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None: # pyright: ignore[reportUnusedFunction]
pass
def test_response_handler_explicit_only_output_raises_error(self):
@@ -858,7 +858,7 @@ class TestResponseHandlerExplicitTypes:
with pytest.raises(ValueError, match="must specify 'request' type"):
@response_handler(output=bool)
- async def test_handler(self, original_request, response, ctx) -> None:
+ async def test_handler(self: Any, original_request: Any, response: Any, ctx: WorkflowContext) -> None: # pyright: ignore[reportUnusedFunction]
pass
def test_executor_with_explicit_response_handlers(self):
@@ -873,7 +873,7 @@ class TestResponseHandlerExplicitTypes:
pass
@response_handler(request=str, response=int, output=bool)
- async def handle_explicit(self, original_request, response, ctx) -> None:
+ async def handle_explicit(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
pass
executor = TestExecutor()
@@ -907,7 +907,7 @@ class TestResponseHandlerExplicitTypes:
pass
@response_handler(request=str, response=int)
- async def handle_response(self, original_request, response, ctx) -> None:
+ async def handle_response(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
self.handled_request = original_request
self.handled_response = response
@@ -942,7 +942,7 @@ class TestResponseHandlerExplicitTypes:
# Explicit type handler
@response_handler(request=dict, response=bool)
- async def handle_explicit(self, original_request, response, ctx) -> None:
+ async def handle_explicit(self, original_request: Any, response: Any, ctx: WorkflowContext) -> None:
pass
executor = TestExecutor()
diff --git a/python/packages/core/tests/workflow/test_runner.py b/python/packages/core/tests/workflow/test_runner.py
index eaf69f90b0..db6dccd9fa 100644
--- a/python/packages/core/tests/workflow/test_runner.py
+++ b/python/packages/core/tests/workflow/test_runner.py
@@ -2,6 +2,7 @@
import asyncio
from dataclasses import dataclass
+from typing import Any
from unittest.mock import AsyncMock, MagicMock
import pytest
@@ -113,7 +114,7 @@ async def test_runner_run_until_convergence():
assert result is not None and result == 10
# iteration count shouldn't be reset after convergence
- assert runner._iteration == 10 # type: ignore
+ assert runner._iteration == 10 # pyright: ignore[reportPrivateUsage]
async def test_runner_run_until_convergence_not_completed():
@@ -173,7 +174,7 @@ async def test_runner_run_iteration_preserves_message_order_per_edge_runner() ->
for index in range(5):
await ctx.send_message(WorkflowMessage(data=MockMessage(data=index), source_id="source"))
- await runner._run_iteration()
+ await runner._run_iteration() # pyright: ignore[reportPrivateUsage]
assert edge_runner.received == [0, 1, 2, 3, 4]
@@ -213,7 +214,7 @@ async def test_runner_run_iteration_delivers_different_edge_runners_concurrently
await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id="source"))
- iteration_task = asyncio.create_task(runner._run_iteration())
+ iteration_task = asyncio.create_task(runner._run_iteration()) # pyright: ignore[reportPrivateUsage]
await blocking_edge_runner.started.wait()
await asyncio.wait_for(probe_edge_runner.probe_completed.wait(), timeout=2.0)
@@ -280,7 +281,7 @@ async def test_fanout_edge_runner_delivers_to_multiple_targets_concurrently() ->
# Queue a message from source (will be delivered to both targets via FanOut)
await ctx.send_message(WorkflowMessage(data=MockMessage(data=1), source_id=source.id))
- iteration_task = asyncio.create_task(runner._run_iteration())
+ iteration_task = asyncio.create_task(runner._run_iteration()) # pyright: ignore[reportPrivateUsage]
# Wait for the blocking executor to start
await blocking_target.started.wait()
@@ -477,11 +478,11 @@ async def test_runner_reset_iteration_count():
ctx = InProcRunnerContext()
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
- runner._iteration = 10
+ runner._iteration = 10 # pyright: ignore[reportPrivateUsage]
runner.reset_iteration_count()
- assert runner._iteration == 0
+ assert runner._iteration == 0 # pyright: ignore[reportPrivateUsage]
class CheckpointingContext(InProcRunnerContext):
@@ -501,18 +502,19 @@ class CheckpointingContext(InProcRunnerContext):
graph_signature_hash: str,
state: State,
previous_checkpoint_id: str | None,
- iteration: int,
+ iteration_count: int,
+ metadata: dict[str, Any] | None = None,
) -> str:
checkpoint = WorkflowCheckpoint(
workflow_name=workflow_name,
graph_signature_hash=graph_signature_hash,
- state=state.export(),
+ state=state.export_state(),
previous_checkpoint_id=previous_checkpoint_id,
- iteration_count=iteration,
+ iteration_count=iteration_count,
)
return await self._storage.save(checkpoint)
- async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None:
+ async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None: # pyright: ignore[reportIncompatibleMethodOverride]
try:
return await self._storage.load(checkpoint_id)
except WorkflowCheckpointException:
@@ -537,7 +539,8 @@ class FailingCheckpointContext(InProcRunnerContext):
graph_signature_hash: str,
state: State,
previous_checkpoint_id: str | None,
- iteration: int,
+ iteration_count: int,
+ metadata: dict[str, Any] | None = None,
) -> str:
raise RuntimeError("Simulated checkpoint failure")
@@ -609,8 +612,8 @@ async def test_runner_restore_from_checkpoint_with_external_storage():
# Restore using external storage
await runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage=storage)
- assert runner._resumed_from_checkpoint is True
- assert runner._iteration == 5
+ assert runner._resumed_from_checkpoint is True # pyright: ignore[reportPrivateUsage]
+ assert runner._iteration == 5 # pyright: ignore[reportPrivateUsage]
assert state.get("test_key") == "test_value"
@@ -684,7 +687,7 @@ async def test_runner_restore_executor_states_invalid_states_type():
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
with pytest.raises(WorkflowCheckpointException, match="not a dictionary"):
- await runner._restore_executor_states()
+ await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage]
async def test_runner_restore_executor_states_invalid_executor_id_type():
@@ -698,7 +701,7 @@ async def test_runner_restore_executor_states_invalid_executor_id_type():
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
with pytest.raises(WorkflowCheckpointException, match="not a string"):
- await runner._restore_executor_states()
+ await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage]
async def test_runner_restore_executor_states_invalid_state_type():
@@ -712,7 +715,7 @@ async def test_runner_restore_executor_states_invalid_state_type():
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
with pytest.raises(WorkflowCheckpointException, match="not a dict"):
- await runner._restore_executor_states()
+ await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage]
async def test_runner_restore_executor_states_invalid_state_keys():
@@ -726,7 +729,7 @@ async def test_runner_restore_executor_states_invalid_state_keys():
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
with pytest.raises(WorkflowCheckpointException, match="not a dict"):
- await runner._restore_executor_states()
+ await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage]
async def test_runner_restore_executor_states_missing_executor():
@@ -739,7 +742,7 @@ async def test_runner_restore_executor_states_missing_executor():
runner = Runner([], {}, state, ctx, "test_name", graph_signature_hash="test_hash")
with pytest.raises(WorkflowCheckpointException, match="not found during state restoration"):
- await runner._restore_executor_states()
+ await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage]
async def test_runner_set_executor_state_invalid_existing_states():
@@ -752,7 +755,7 @@ async def test_runner_set_executor_state_invalid_existing_states():
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
with pytest.raises(WorkflowCheckpointException, match="not a dictionary"):
- await runner._set_executor_state("executor_a", {"key": "value"})
+ await runner._set_executor_state("executor_a", {"key": "value"}) # pyright: ignore[reportPrivateUsage]
async def test_runner_with_pre_loop_events():
@@ -779,7 +782,7 @@ class EventEmittingExecutor(Executor):
"""An executor that emits events during execution."""
@handler
- async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+ async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, str]) -> None:
# Emit event during processing
await ctx.yield_output(f"processed-{message.data}")
if message.data < 3:
@@ -831,7 +834,7 @@ async def test_runner_restore_executor_states_no_states():
runner = Runner([], {executor_a.id: executor_a}, state, ctx, "test_name", graph_signature_hash="test_hash")
# Should complete without error when no executor states exist
- await runner._restore_executor_states()
+ await runner._restore_executor_states() # pyright: ignore[reportPrivateUsage]
async def test_runner_checkpoint_with_resumed_flag():
@@ -853,7 +856,7 @@ async def test_runner_checkpoint_with_resumed_flag():
state = State()
runner = Runner(edges, executors, state, ctx, "test_name", graph_signature_hash="test_hash")
- runner._mark_resumed(5)
+ runner._mark_resumed(5) # pyright: ignore[reportPrivateUsage]
# Add a message to trigger the checkpoint creation path
await ctx.send_message(WorkflowMessage(data=MockMessage(data=8), source_id="START"))
@@ -870,7 +873,7 @@ async def test_runner_checkpoint_with_resumed_flag():
pass
# After completing, resumed flag should be reset
- assert runner._resumed_from_checkpoint is False
+ assert runner._resumed_from_checkpoint is False # pyright: ignore[reportPrivateUsage]
class ExecutorThatFailsWithEvents(Executor):
@@ -883,7 +886,7 @@ class ExecutorThatFailsWithEvents(Executor):
self._iteration_count = 0
@handler
- async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+ async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, str]) -> None:
self._iteration_count += 1
# First emit an output event to the workflow context
await ctx.yield_output(f"output-before-failure-{message.data}")
@@ -951,7 +954,7 @@ class SlowEventEmittingExecutor(Executor):
self.current_iteration = 0
@handler
- async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, int]) -> None:
+ async def handle(self, message: MockMessage, ctx: WorkflowContext[MockMessage, str]) -> None:
self.current_iteration += 1
# Emit output event
await ctx.yield_output(f"iteration-{self.current_iteration}")
diff --git a/python/packages/core/tests/workflow/test_state.py b/python/packages/core/tests/workflow/test_state.py
index 486fc9fa25..7781eb4141 100644
--- a/python/packages/core/tests/workflow/test_state.py
+++ b/python/packages/core/tests/workflow/test_state.py
@@ -61,9 +61,9 @@ class TestSuperstepCaching:
state.set("key", "value")
# Value is in pending
- assert "key" in state._pending
+ assert "key" in state._pending # pyright: ignore[reportPrivateUsage]
# Value is NOT in committed
- assert "key" not in state._committed
+ assert "key" not in state._committed # pyright: ignore[reportPrivateUsage]
# But get() still returns it
assert state.get("key") == "value"
@@ -72,14 +72,14 @@ class TestSuperstepCaching:
state.set("key", "value")
# Before commit: in pending, not committed
- assert "key" in state._pending
- assert "key" not in state._committed
+ assert "key" in state._pending # pyright: ignore[reportPrivateUsage]
+ assert "key" not in state._committed # pyright: ignore[reportPrivateUsage]
state.commit()
# After commit: in committed, pending cleared
- assert "key" not in state._pending
- assert "key" in state._committed
+ assert "key" not in state._pending # pyright: ignore[reportPrivateUsage]
+ assert "key" in state._committed # pyright: ignore[reportPrivateUsage]
assert state.get("key") == "value"
def test_discard_clears_pending_without_committing(self) -> None:
@@ -108,7 +108,7 @@ class TestSuperstepCaching:
# get() returns pending value, not committed
assert state.get("key") == "pending_value"
# But committed still has old value
- assert state._committed["key"] == "committed_value"
+ assert state._committed["key"] == "committed_value" # pyright: ignore[reportPrivateUsage]
def test_multiple_sets_before_commit(self) -> None:
state = State()
@@ -130,13 +130,13 @@ class TestDeleteWithSuperstepCaching:
state = State()
state.set("key", "value")
# Key only in pending, not committed
- assert "key" in state._pending
- assert "key" not in state._committed
+ assert "key" in state._pending # pyright: ignore[reportPrivateUsage]
+ assert "key" not in state._committed # pyright: ignore[reportPrivateUsage]
state.delete("key")
# Should be removed from pending
- assert "key" not in state._pending
+ assert "key" not in state._pending # pyright: ignore[reportPrivateUsage]
assert state.get("key") is None
assert state.has("key") is False
@@ -148,14 +148,14 @@ class TestDeleteWithSuperstepCaching:
state.delete("key")
# Key should be marked for deletion in pending (sentinel)
- assert "key" in state._pending
+ assert "key" in state._pending # pyright: ignore[reportPrivateUsage]
# get() should return default (not the sentinel!)
assert state.get("key") is None
assert state.get("key", "default") == "default"
# has() should return False
assert state.has("key") is False
# But committed still has it until commit()
- assert "key" in state._committed
+ assert "key" in state._committed # pyright: ignore[reportPrivateUsage]
def test_delete_committed_key_removed_on_commit(self) -> None:
state = State()
@@ -166,8 +166,8 @@ class TestDeleteWithSuperstepCaching:
state.commit()
# Now it should be gone from committed too
- assert "key" not in state._committed
- assert "key" not in state._pending
+ assert "key" not in state._committed # pyright: ignore[reportPrivateUsage]
+ assert "key" not in state._pending # pyright: ignore[reportPrivateUsage]
def test_delete_key_in_both_pending_and_committed(self) -> None:
"""Test delete when key exists in both pending (modified) and committed."""
@@ -177,8 +177,8 @@ class TestDeleteWithSuperstepCaching:
# Modify the key (now in both pending and committed)
state.set("key", "modified")
- assert state._pending["key"] == "modified"
- assert state._committed["key"] == "original"
+ assert state._pending["key"] == "modified" # pyright: ignore[reportPrivateUsage]
+ assert state._committed["key"] == "original" # pyright: ignore[reportPrivateUsage]
# Delete should mark for deletion from committed
state.delete("key")
@@ -189,8 +189,8 @@ class TestDeleteWithSuperstepCaching:
# After commit, key should be fully removed
state.commit()
- assert "key" not in state._committed
- assert "key" not in state._pending
+ assert "key" not in state._committed # pyright: ignore[reportPrivateUsage]
+ assert "key" not in state._pending # pyright: ignore[reportPrivateUsage]
def test_discard_after_delete_restores_committed_value(self) -> None:
state = State()
@@ -238,12 +238,12 @@ class TestFailureScenarios:
state.set("key3", "value3")
# Before commit - nothing in committed
- assert len(state._committed) == 0
+ assert len(state._committed) == 0 # pyright: ignore[reportPrivateUsage]
state.commit()
# After commit - all three values committed together
- assert state._committed == {"key1": "value1", "key2": "value2", "key3": "value3"}
+ assert state._committed == {"key1": "value1", "key2": "value2", "key3": "value3"} # pyright: ignore[reportPrivateUsage]
def test_repeated_supersteps_are_isolated(self) -> None:
"""Test that each superstep's changes are isolated until committed."""
@@ -300,4 +300,4 @@ class TestExportImport:
# Pending is still there
assert state.get("pending_key") == "pending_value"
- assert "pending_key" in state._pending
+ assert "pending_key" in state._pending # pyright: ignore[reportPrivateUsage]
diff --git a/python/packages/core/tests/workflow/test_typing_utils.py b/python/packages/core/tests/workflow/test_typing_utils.py
index 4dc8d8c917..f94bd9d52e 100644
--- a/python/packages/core/tests/workflow/test_typing_utils.py
+++ b/python/packages/core/tests/workflow/test_typing_utils.py
@@ -36,32 +36,32 @@ def test_normalize_type_to_list_none() -> None:
def test_normalize_type_to_list_union_pipe_syntax() -> None:
"""Test normalize_type_to_list with union types using | syntax."""
- result = normalize_type_to_list(str | int)
+ result = normalize_type_to_list(str | int) # pyright: ignore[reportArgumentType]
assert set(result) == {str, int}
- result = normalize_type_to_list(str | int | bool)
+ result = normalize_type_to_list(str | int | bool) # pyright: ignore[reportArgumentType]
assert set(result) == {str, int, bool}
def test_normalize_type_to_list_union_typing_syntax() -> None:
"""Test normalize_type_to_list with Union[] from typing module."""
- result = normalize_type_to_list(Union[str, int])
+ result = normalize_type_to_list(Union[str, int]) # pyright: ignore[reportArgumentType]
assert set(result) == {str, int}
- result = normalize_type_to_list(Union[str, int, bool])
+ result = normalize_type_to_list(Union[str, int, bool]) # pyright: ignore[reportArgumentType]
assert set(result) == {str, int, bool}
def test_normalize_type_to_list_optional() -> None:
"""Test normalize_type_to_list with Optional types (Union[T, None])."""
# Optional[str] is Union[str, None]
- result = normalize_type_to_list(Optional[str])
+ result = normalize_type_to_list(Optional[str]) # pyright: ignore[reportArgumentType]
assert str in result
assert type(None) in result
assert len(result) == 2
# str | None is equivalent
- result = normalize_type_to_list(str | None)
+ result = normalize_type_to_list(str | None) # pyright: ignore[reportArgumentType]
assert str in result
assert type(None) in result
assert len(result) == 2
@@ -77,7 +77,7 @@ def test_normalize_type_to_list_custom_types() -> None:
result = normalize_type_to_list(CustomMessage)
assert result == [CustomMessage]
- result = normalize_type_to_list(CustomMessage | str)
+ result = normalize_type_to_list(CustomMessage | str) # pyright: ignore[reportArgumentType]
assert set(result) == {CustomMessage, str}
@@ -96,7 +96,7 @@ def test_resolve_type_annotation_actual_types() -> None:
"""Test resolve_type_annotation passes through actual types unchanged."""
assert resolve_type_annotation(str) is str
assert resolve_type_annotation(int) is int
- assert resolve_type_annotation(str | int) == str | int
+ assert resolve_type_annotation(str | int) == str | int # pyright: ignore[reportArgumentType]
def test_resolve_type_annotation_string_builtin() -> None:
diff --git a/python/packages/core/tests/workflow/test_validation.py b/python/packages/core/tests/workflow/test_validation.py
index ae694c8354..be3c8b45f7 100644
--- a/python/packages/core/tests/workflow/test_validation.py
+++ b/python/packages/core/tests/workflow/test_validation.py
@@ -484,8 +484,8 @@ def test_handler_ctx_missing_annotation_raises() -> None:
# Validation now happens at handler registration time, not workflow build time
with pytest.raises(ValueError) as exc:
- class BadExecutor(Executor):
- @handler
+ class BadExecutor(Executor): # pyright: ignore[reportUnusedClass]
+ @handler # pyright: ignore[reportUnknownArgumentType]
async def handle(self, message: str, ctx) -> None: # type: ignore[no-untyped-def]
pass
@@ -496,8 +496,8 @@ def test_handler_ctx_invalid_t_out_entries_raises() -> None:
# Validation now happens at handler registration time, not workflow build time
with pytest.raises(ValueError) as exc:
- class BadExecutor(Executor):
- @handler
+ class BadExecutor(Executor): # pyright: ignore[reportUnusedClass]
+ @handler # pyright: ignore[reportUnknownArgumentType]
async def handle(self, message: str, ctx: WorkflowContext[123]) -> None: # type: ignore[valid-type]
pass
@@ -555,7 +555,7 @@ def test_output_validation_with_valid_output_executors():
)
assert workflow is not None
- assert workflow._output_executors == ["executor2"]
+ assert workflow._output_executors == ["executor2"] # pyright: ignore[reportPrivateUsage]
def test_output_validation_with_multiple_valid_output_executors():
@@ -572,7 +572,7 @@ def test_output_validation_with_multiple_valid_output_executors():
)
assert workflow is not None
- assert set(workflow._output_executors) == {"executor1", "executor3"}
+ assert set(workflow._output_executors) == {"executor1", "executor3"} # pyright: ignore[reportPrivateUsage]
def test_output_validation_fails_for_nonexistent_executor():
diff --git a/python/packages/core/tests/workflow/test_viz.py b/python/packages/core/tests/workflow/test_viz.py
index bf7bbffee1..5573dadd61 100644
--- a/python/packages/core/tests/workflow/test_viz.py
+++ b/python/packages/core/tests/workflow/test_viz.py
@@ -2,6 +2,9 @@
"""Tests for the workflow visualization module."""
+from pathlib import Path
+from typing import Any
+
import pytest
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, WorkflowExecutor, WorkflowViz, handler
@@ -25,7 +28,7 @@ class ListStrTargetExecutor(Executor):
@pytest.fixture
-def basic_sub_workflow():
+def basic_sub_workflow() -> dict[str, Any]:
"""Fixture that creates a basic sub-workflow setup for testing."""
# Create a sub-workflow
sub_exec1 = MockExecutor(id="sub_exec1")
@@ -98,7 +101,7 @@ def test_workflow_viz_export_dot():
assert '"executor1" -> "executor2"' in content
-def test_workflow_viz_export_dot_with_filename(tmp_path):
+def test_workflow_viz_export_dot_with_filename(tmp_path: Path):
"""Test exporting workflow as DOT format with specified filename."""
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
@@ -203,7 +206,7 @@ def test_workflow_viz_graphviz_binary_not_found():
mock_source_class.return_value = mock_source
# Import the ExecutableNotFound exception for the test
- from graphviz.backend.execute import ExecutableNotFound
+ from graphviz.backend.execute import ExecutableNotFound # type: ignore[import-not-found]
mock_source.render.side_effect = ExecutableNotFound("failed to execute PosixPath('dot')")
@@ -329,7 +332,7 @@ def test_workflow_viz_mermaid_fan_in_edge_group():
assert "s2 --> t" not in mermaid
-def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow):
+def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow: dict[str, Any]):
"""Test that WorkflowViz can visualize sub-workflows in DOT format."""
main_workflow = basic_sub_workflow["main_workflow"]
@@ -353,7 +356,7 @@ def test_workflow_viz_sub_workflow_digraph(basic_sub_workflow):
assert '"workflow_executor_1/sub_exec1" -> "workflow_executor_1/sub_exec2"' in dot_content
-def test_workflow_viz_sub_workflow_mermaid(basic_sub_workflow):
+def test_workflow_viz_sub_workflow_mermaid(basic_sub_workflow: dict[str, Any]):
"""Test that WorkflowViz can visualize sub-workflows in Mermaid format."""
main_workflow = basic_sub_workflow["main_workflow"]
diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py
index 8bbf11fa6a..f338ce94f6 100644
--- a/python/packages/core/tests/workflow/test_workflow.py
+++ b/python/packages/core/tests/workflow/test_workflow.py
@@ -4,7 +4,7 @@ import asyncio
import tempfile
from collections.abc import AsyncIterable, Awaitable, Sequence
from dataclasses import dataclass, field
-from typing import Any, cast
+from typing import Any, Literal, cast, overload
from uuid import uuid4
import pytest
@@ -13,6 +13,7 @@ from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
+ AgentRunInputs,
AgentSession,
BaseAgent,
Content,
@@ -474,7 +475,7 @@ class StateTrackingExecutor(Executor):
) -> None:
"""Handle the message and track it in workflow state."""
# Get existing messages from workflow state
- existing_messages = ctx.get_state("processed_messages") or []
+ existing_messages: list[str] = ctx.get_state("processed_messages") or []
# Record this message
message_record = f"{message.run_id}:{message.data}"
@@ -833,6 +834,26 @@ class _StreamingTestAgent(BaseAgent):
super().__init__(**kwargs)
self._reply_text = reply_text
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
@@ -883,8 +904,10 @@ async def test_agent_streaming_vs_non_streaming() -> None:
stream_events.append(event)
# Filter for agent events
- agent_response = [
- cast(AgentResponse, e.data) for e in stream_events if e.type == "output" and isinstance(e.data, AgentResponse)
+ agent_response: list[AgentResponse[Any]] = [
+ cast(AgentResponse[Any], e.data) # pyright: ignore[reportUnknownMemberType]
+ for e in stream_events
+ if e.type == "output" and isinstance(e.data, AgentResponse)
]
agent_response_updates = [
e.data for e in stream_events if e.type == "output" and isinstance(e.data, AgentResponseUpdate)
diff --git a/python/packages/core/tests/workflow/test_workflow_agent.py b/python/packages/core/tests/workflow/test_workflow_agent.py
index d20d60ba3b..b5a8bb9902 100644
--- a/python/packages/core/tests/workflow/test_workflow_agent.py
+++ b/python/packages/core/tests/workflow/test_workflow_agent.py
@@ -2,7 +2,7 @@
import uuid
from collections.abc import Awaitable, Sequence
-from typing import Any
+from typing import Any, Literal, overload
import pytest
from typing_extensions import Never
@@ -713,6 +713,14 @@ class TestWorkflowAgent:
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
+ def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
+ return AgentSession()
+
+ @overload
+ def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
@@ -801,6 +809,14 @@ class TestWorkflowAgent:
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
+ def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
+ return AgentSession()
+
+ @overload
+ def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: str | Content | Message | Sequence[str | Content | Message] | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
@@ -1207,7 +1223,7 @@ class TestWorkflowAgentMergeUpdates:
]
# Compare using role.value for Role enum
- actual_sequence_normalized = [(t, r.value if hasattr(r, "value") else r) for t, r in content_sequence]
+ actual_sequence_normalized = [(t, r.value if hasattr(r, "value") else r) for t, r in content_sequence] # type: ignore[union-attr]
assert actual_sequence_normalized == expected_sequence, (
f"FunctionResultContent should come immediately after FunctionCallContent. "
diff --git a/python/packages/core/tests/workflow/test_workflow_builder.py b/python/packages/core/tests/workflow/test_workflow_builder.py
index 073a24e5a3..3a7b719530 100644
--- a/python/packages/core/tests/workflow/test_workflow_builder.py
+++ b/python/packages/core/tests/workflow/test_workflow_builder.py
@@ -1,7 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.
+from collections.abc import AsyncIterator, Awaitable
from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal, overload
import pytest
@@ -9,10 +10,12 @@ from agent_framework import (
AgentExecutor,
AgentResponse,
AgentResponseUpdate,
+ AgentRunInputs,
AgentSession,
BaseAgent,
Executor,
Message,
+ ResponseStream,
WorkflowBuilder,
WorkflowContext,
WorkflowValidationError,
@@ -21,22 +24,49 @@ from agent_framework import (
class DummyAgent(BaseAgent):
- def run(self, messages=None, *, stream: bool = False, session: AgentSession | None = None, **kwargs): # type: ignore[override]
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[False] = ...,
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]]: ...
+
+ @overload
+ def run(
+ self,
+ messages: AgentRunInputs | None = ...,
+ *,
+ stream: Literal[True],
+ session: AgentSession | None = ...,
+ **kwargs: Any,
+ ) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
+ def run(
+ self,
+ messages: AgentRunInputs | None = None,
+ *,
+ stream: bool = False,
+ session: AgentSession | None = None,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
if stream:
- return self._run_stream_impl()
+ return ResponseStream[AgentResponseUpdate, AgentResponse[Any]](self._run_stream_impl())
return self._run_impl(messages)
- async def _run_impl(self, messages=None) -> AgentResponse:
+ async def _run_impl(self, messages: AgentRunInputs | None = None) -> AgentResponse:
norm: list[Message] = []
if messages:
- for m in messages: # type: ignore[iteration-over-optional]
+ for m in messages: # type: ignore[union-attr]
if isinstance(m, Message):
norm.append(m)
elif isinstance(m, str):
norm.append(Message(role="user", text=m))
return AgentResponse(messages=norm)
- async def _run_stream_impl(self): # type: ignore[override]
+ async def _run_stream_impl(self) -> AsyncIterator[AgentResponseUpdate]:
# Minimal async generator
yield AgentResponseUpdate()
@@ -202,7 +232,7 @@ def test_with_output_from_returns_builder():
builder = WorkflowBuilder(output_executors=[executor_a], start_executor=executor_a)
# Verify builder was created with output_executors
- assert builder._output_executors == [executor_a]
+ assert builder._output_executors == [executor_a] # pyright: ignore[reportPrivateUsage]
def test_with_output_from_with_executor_instances():
diff --git a/python/packages/core/tests/workflow/test_workflow_context.py b/python/packages/core/tests/workflow/test_workflow_context.py
index 53a7e44903..a13c0b5a55 100644
--- a/python/packages/core/tests/workflow/test_workflow_context.py
+++ b/python/packages/core/tests/workflow/test_workflow_context.py
@@ -84,7 +84,7 @@ async def test_executor_emits_normal_event() -> None:
class _TestEvent(WorkflowEvent):
def __init__(self, data: Any = None) -> None:
- super().__init__("test_event", data=data)
+ super().__init__("test_event", data=data) # type: ignore[arg-type]
async def test_workflow_context_type_annotations_no_parameter() -> None:
@@ -244,8 +244,8 @@ async def test_workflow_context_missing_annotation_error() -> None:
# Test class-based executor with missing ctx annotation
with pytest.raises(ValueError, match="must have a WorkflowContext"):
- class _BadExecutor(Executor):
- @handler
+ class _BadExecutor(Executor): # pyright: ignore[reportUnusedClass]
+ @handler # pyright: ignore[reportUnknownArgumentType]
async def bad_handler(self, text: str, ctx) -> None: # type: ignore[no-untyped-def]
pass
@@ -264,8 +264,8 @@ async def test_workflow_context_invalid_type_parameter_error() -> None:
# Test class-based executor with invalid type parameter
with pytest.raises(ValueError, match="invalid type entry"):
- class _BadExecutor(Executor):
- @handler
+ class _BadExecutor(Executor): # pyright: ignore[reportUnusedClass]
+ @handler # pyright: ignore[reportUnknownArgumentType]
async def bad_handler(self, text: str, ctx: WorkflowContext[456]) -> None: # type: ignore[valid-type]
pass
diff --git a/python/packages/core/tests/workflow/test_workflow_kwargs.py b/python/packages/core/tests/workflow/test_workflow_kwargs.py
index ce1465effc..0850c6b060 100644
--- a/python/packages/core/tests/workflow/test_workflow_kwargs.py
+++ b/python/packages/core/tests/workflow/test_workflow_kwargs.py
@@ -1,13 +1,14 @@
# Copyright (c) Microsoft. All rights reserved.
-from collections.abc import AsyncIterable, Awaitable, Sequence
-from typing import Annotated, Any
+from collections.abc import AsyncIterable, Awaitable
+from typing import Annotated, Any, Literal, overload
import pytest
from agent_framework import (
AgentResponse,
AgentResponseUpdate,
+ AgentRunInputs,
AgentSession,
BaseAgent,
Content,
@@ -50,14 +51,19 @@ class _KwargsCapturingAgent(BaseAgent):
super().__init__(name=name, description="Test agent for kwargs capture")
self.captured_kwargs = []
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.captured_kwargs.append(dict(kwargs))
if stream:
@@ -83,15 +89,20 @@ class _OptionsAwareAgent(BaseAgent):
self.captured_options = []
self.captured_kwargs = []
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
options: dict[str, Any] | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.captured_options.append(dict(options) if options is not None else None)
self.captured_kwargs.append(dict(kwargs))
if stream:
@@ -189,15 +200,15 @@ async def test_sequential_run_options_does_not_conflict_with_agent_options() ->
break
assert len(agent.captured_options) >= 1
- captured_options = agent.captured_options[0]
+ captured_options: dict[str, Any] | None = agent.captured_options[0]
assert captured_options is not None
assert captured_options.get("store") is False
- additional_args = captured_options.get("additional_function_arguments")
+ additional_args: Any = captured_options.get("additional_function_arguments")
assert isinstance(additional_args, dict)
- assert additional_args.get("source") == "workflow-options"
- assert additional_args.get("custom_data") == custom_data
- assert additional_args.get("user_token") == user_token
+ assert additional_args.get("source") == "workflow-options" # pyright: ignore[reportUnknownMemberType]
+ assert additional_args.get("custom_data") == custom_data # pyright: ignore[reportUnknownMemberType]
+ assert additional_args.get("user_token") == user_token # pyright: ignore[reportUnknownMemberType]
# "options" should be passed once via the dedicated options parameter,
# not duplicated in **kwargs.
@@ -225,13 +236,13 @@ async def test_sequential_run_additional_function_arguments_flattened() -> None:
break
assert len(agent.captured_options) >= 1
- captured_options = agent.captured_options[0]
+ captured_options: dict[str, Any] | None = agent.captured_options[0]
assert captured_options is not None
- additional_args = captured_options.get("additional_function_arguments")
+ additional_args: Any = captured_options.get("additional_function_arguments")
assert isinstance(additional_args, dict)
- assert additional_args.get("custom_data") == custom_data
- assert additional_args.get("user_token") == user_token
+ assert additional_args.get("custom_data") == custom_data # pyright: ignore[reportUnknownMemberType]
+ assert additional_args.get("user_token") == user_token # pyright: ignore[reportUnknownMemberType]
assert "additional_function_arguments" not in additional_args
assert len(agent.captured_kwargs) >= 1
@@ -255,14 +266,14 @@ async def test_sequential_run_additional_function_arguments_merges_with_options(
break
assert len(agent.captured_options) >= 1
- captured_options = agent.captured_options[0]
+ captured_options: dict[str, Any] | None = agent.captured_options[0]
assert captured_options is not None
- additional_args = captured_options.get("additional_function_arguments")
+ additional_args: Any = captured_options.get("additional_function_arguments")
assert isinstance(additional_args, dict)
- assert additional_args.get("source") == "workflow-options"
- assert additional_args.get("custom_data") == {"session_id": "abc123"}
- assert additional_args.get("user_token") == {"user_name": "alice"}
+ assert additional_args.get("source") == "workflow-options" # pyright: ignore[reportUnknownMemberType]
+ assert additional_args.get("custom_data") == {"session_id": "abc123"} # pyright: ignore[reportUnknownMemberType]
+ assert additional_args.get("user_token") == {"user_name": "alice"} # pyright: ignore[reportUnknownMemberType]
assert "additional_function_arguments" not in additional_args
@@ -463,14 +474,19 @@ async def test_kwargs_preserved_on_response_continuation() -> None:
self.captured_kwargs = []
self._asked = False
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.captured_kwargs.append(dict(kwargs))
if not self._asked:
self._asked = True
@@ -521,14 +537,19 @@ async def test_kwargs_overridden_on_response_continuation() -> None:
self.captured_kwargs = []
self._asked = False
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.captured_kwargs.append(dict(kwargs))
if not self._asked:
self._asked = True
@@ -583,14 +604,19 @@ async def test_kwargs_empty_value_passed_on_continuation() -> None:
self.captured_kwargs = []
self._asked = False
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[False] = ..., session: AgentSession | None = ..., **kwargs: Any) -> Awaitable[AgentResponse[Any]]: ...
+ @overload
+ def run(self, messages: AgentRunInputs | None = ..., *, stream: Literal[True], session: AgentSession | None = ..., **kwargs: Any) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
+
def run(
self,
- messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
- ) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
+ ) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.captured_kwargs.append(dict(kwargs))
if not self._asked:
self._asked = True
@@ -690,8 +716,8 @@ async def test_handoff_kwargs_flow_to_agents() -> None:
workflow = (
HandoffBuilder(termination_condition=lambda conv: len(conv) >= 4)
- .participants([agent1, agent2])
- .with_start_agent(agent1)
+ .participants([agent1, agent2]) # type: ignore[list-item]
+ .with_start_agent(agent1) # type: ignore[arg-type]
.with_autonomous_mode()
.build()
)
diff --git a/python/packages/core/tests/workflow/test_workflow_observability.py b/python/packages/core/tests/workflow/test_workflow_observability.py
index b2260abe63..b098fa2771 100644
--- a/python/packages/core/tests/workflow/test_workflow_observability.py
+++ b/python/packages/core/tests/workflow/test_workflow_observability.py
@@ -109,7 +109,7 @@ async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter)
{
"id": "test-workflow-123",
"max_iterations": 100,
- "model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}',
+ "model_dump_json": lambda self: '{"id": "test-workflow-123", "type": "mock"}', # pyright: ignore[reportUnknownLambdaType]
},
)(),
)
@@ -122,7 +122,7 @@ async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter)
},
) as workflow_span:
workflow_span.add_event(OtelAttr.WORKFLOW_STARTED)
- sending_attributes = {
+ sending_attributes: dict[str, str | int] = {
OtelAttr.MESSAGE_TYPE: "ResponseMessage",
OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID: "target-789",
}
@@ -231,7 +231,7 @@ async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> No
@pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
async def test_trace_context_disabled_when_tracing_disabled(
- enable_instrumentation, span_exporter: InMemorySpanExporter
+ enable_instrumentation: bool, span_exporter: InMemorySpanExporter
) -> None:
"""Test that no trace context is added when tracing is disabled."""
# Tracing should be disabled by default
@@ -313,7 +313,7 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
span_exporter.clear()
# Run workflow (this should create run spans)
- events = []
+ events: list[Any] = []
async for event in workflow.run("test input", stream=True):
events.append(event)
diff --git a/python/packages/core/tests/workflow/test_workflow_states.py b/python/packages/core/tests/workflow/test_workflow_states.py
index 0ccf84b103..34c7e8c93f 100644
--- a/python/packages/core/tests/workflow/test_workflow_states.py
+++ b/python/packages/core/tests/workflow/test_workflow_states.py
@@ -1,5 +1,7 @@
# Copyright (c) Microsoft. All rights reserved.
+from typing import Any
+
import pytest
from typing_extensions import Never
@@ -36,16 +38,16 @@ async def test_executor_failed_and_workflow_failed_events_streaming():
events.append(ev)
# executor_failed event (type='executor_failed') should be emitted before workflow failed event
- executor_failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+ executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
assert executor_failed_events, "executor_failed event should be emitted when start executor fails"
assert executor_failed_events[0].executor_id == "f"
assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
# Workflow-level failure and FAILED status should be surfaced
- failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
+ failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
assert failed_events
assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
- status = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
+ status: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
assert status and status[-1].state == WorkflowRunState.FAILED
assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in status)
@@ -94,13 +96,13 @@ async def test_executor_failed_event_from_second_executor_in_chain():
events.append(ev)
# executor_failed event should be emitted for the failing executor
- executor_failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
+ executor_failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_failed"]
assert executor_failed_events, "executor_failed event should be emitted when second executor fails"
assert executor_failed_events[0].executor_id == "failing"
assert executor_failed_events[0].origin is WorkflowEventSource.FRAMEWORK
# Workflow-level failure should also be surfaced
- failed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
+ failed_events: list[WorkflowEvent[Any]] = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "failed"]
assert failed_events
assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed_events)
diff --git a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py
index 687dad096b..01a68e6a8e 100644
--- a/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py
+++ b/python/packages/declarative/agent_framework_declarative/_workflows/_declarative_base.py
@@ -388,11 +388,15 @@ class DeclarativeWorkflowState:
from System.Globalization import CultureInfo
original_culture = CultureInfo.CurrentCulture
- CultureInfo.CurrentCulture = CultureInfo("en-US")
+ original_ui_culture = CultureInfo.CurrentUICulture
+ en_us_culture = CultureInfo("en-US")
+ CultureInfo.CurrentCulture = en_us_culture
+ CultureInfo.CurrentUICulture = en_us_culture
try:
return engine.eval(formula, symbols=symbols)
finally:
CultureInfo.CurrentCulture = original_culture
+ CultureInfo.CurrentUICulture = original_ui_culture
except ValueError as e:
error_msg = str(e)
# Handle undefined variable errors gracefully by returning None
diff --git a/python/packages/declarative/tests/test_powerfx_yaml_compatibility.py b/python/packages/declarative/tests/test_powerfx_yaml_compatibility.py
index 9591dc05cb..8ea3c3af57 100644
--- a/python/packages/declarative/tests/test_powerfx_yaml_compatibility.py
+++ b/python/packages/declarative/tests/test_powerfx_yaml_compatibility.py
@@ -493,6 +493,31 @@ class TestPowerFxUndefinedVariables:
result = state.eval("=Local.Something.Nested.Deep")
assert result is None
+ async def test_undefined_variable_returns_none_with_non_english_ui_culture(self, mock_state):
+ """Test that undefined variables return None even when CurrentUICulture is non-English.
+
+ Regression test for #4321: on non-English systems, CurrentUICulture causes
+ PowerFx to emit localized error messages that don't match the English
+ string guards ("isn't recognized", "Name isn't valid"), crashing the workflow.
+ The fix sets CurrentUICulture to en-US alongside CurrentCulture before eval.
+ """
+ from System.Globalization import CultureInfo
+
+ state = DeclarativeWorkflowState(mock_state)
+ state.initialize()
+
+ # Simulate a non-English UI culture (e.g. Italian)
+ original_ui_culture = CultureInfo.CurrentUICulture
+ CultureInfo.CurrentUICulture = CultureInfo("it-IT")
+ try:
+ # Should return None, not raise ValueError with Italian error text
+ result = state.eval("=Local.StatusConversationId")
+ assert result is None
+ # Verify the production code restored CurrentUICulture after eval
+ assert str(CultureInfo.CurrentUICulture) == str(CultureInfo("it-IT"))
+ finally:
+ CultureInfo.CurrentUICulture = original_ui_culture
+
class TestStringInterpolation:
"""Test string interpolation patterns."""
diff --git a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py
index 17b927326b..b887d86df3 100644
--- a/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py
+++ b/python/packages/orchestrations/agent_framework_orchestrations/_magentic.py
@@ -14,6 +14,7 @@ from typing import Any, ClassVar, TypeVar, cast
from agent_framework import (
AgentResponse,
+ AgentSession,
Message,
SupportsAgentRun,
)
@@ -559,6 +560,7 @@ class StandardMagenticManager(MagenticManagerBase):
)
self._agent: SupportsAgentRun = agent
+ self._session: AgentSession = self._agent.create_session()
self.task_ledger: _MagenticTaskLedger | None = task_ledger
# Prompts may be overridden if needed
@@ -587,7 +589,7 @@ class StandardMagenticManager(MagenticManagerBase):
The agent's run method is called which applies the agent's configured options
(temperature, seed, instructions, etc.).
"""
- response: AgentResponse = await self._agent.run(messages)
+ response: AgentResponse = await self._agent.run(messages, session=self._session)
if not response.messages:
raise RuntimeError("Agent returned no messages in response.")
if len(response.messages) > 1:
@@ -730,6 +732,7 @@ class StandardMagenticManager(MagenticManagerBase):
state: dict[str, Any] = {}
if self.task_ledger is not None:
state["task_ledger"] = self.task_ledger.to_dict()
+ state["agent_session"] = self._session.to_dict()
return state
@override
@@ -740,6 +743,12 @@ class StandardMagenticManager(MagenticManagerBase):
self.task_ledger = _MagenticTaskLedger.from_dict(ledger)
except Exception: # pragma: no cover - defensive
logger.warning("Failed to restore manager task ledger from checkpoint state")
+ session_payload = state.get("agent_session")
+ if session_payload is not None:
+ try:
+ self._session = AgentSession.from_dict(session_payload)
+ except Exception: # pragma: no cover - defensive
+ logger.warning("Failed to restore manager agent session from checkpoint state")
# endregion Magentic Manager
diff --git a/python/packages/orchestrations/tests/test_magentic.py b/python/packages/orchestrations/tests/test_magentic.py
index e1d0ef8c32..1857a16ee4 100644
--- a/python/packages/orchestrations/tests/test_magentic.py
+++ b/python/packages/orchestrations/tests/test_magentic.py
@@ -1074,4 +1074,71 @@ def test_magentic_agent_factory_with_standard_manager_options():
assert manager.final_answer_prompt == custom_final_prompt
+async def test_standard_manager_propagates_session_to_agent():
+ """Verify StandardMagenticManager passes a consistent session to the underlying agent.
+
+ Regression test for #4371: context providers (e.g. RedisHistoryProvider) configured on
+ the manager agent silently failed because no session was propagated.
+ """
+ captured_sessions: list[AgentSession | None] = []
+
+ class SessionCapturingAgent(BaseAgent):
+ """Agent that records the session passed to each run() call."""
+
+ def run(
+ self,
+ messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
+ *,
+ stream: bool = False,
+ session: Any = None,
+ **kwargs: Any,
+ ) -> Awaitable[AgentResponse] | AsyncIterable[AgentResponseUpdate]:
+ captured_sessions.append(session)
+
+ async def _run() -> AgentResponse:
+ return AgentResponse(messages=[Message("assistant", ["ok"])])
+
+ return _run()
+
+ agent = SessionCapturingAgent()
+ mgr = StandardMagenticManager(agent=agent)
+ ctx = MagenticContext(task="task", participant_descriptions={"a": "desc"})
+
+ await mgr.plan(ctx.clone())
+
+ # plan() calls _complete twice (facts + plan), both should receive the same session
+ assert len(captured_sessions) == 2
+ assert all(s is not None for s in captured_sessions), "session must be passed to agent.run()"
+ assert captured_sessions[0] is captured_sessions[1], "same session instance must be reused across calls"
+ assert captured_sessions[0] is mgr._session
+
+
+def test_standard_manager_checkpoint_preserves_session():
+ """Verify that checkpoint save/restore preserves the manager's session identity."""
+ agent = StubManagerAgent()
+ mgr = StandardMagenticManager(agent=agent)
+ original_session_id = mgr._session.session_id
+
+ state = mgr.on_checkpoint_save()
+ assert "agent_session" in state
+
+ # Restore into a fresh manager and verify session_id is preserved
+ mgr2 = StandardMagenticManager(agent=agent)
+ assert mgr2._session.session_id != original_session_id
+ mgr2.on_checkpoint_restore(state)
+ assert mgr2._session.session_id == original_session_id
+
+
+def test_standard_manager_checkpoint_restore_empty_state():
+ """Verify that restoring from a state without agent_session leaves the session intact."""
+ agent = StubManagerAgent()
+ mgr = StandardMagenticManager(agent=agent)
+ original_session = mgr._session
+ original_session_id = original_session.session_id
+
+ mgr.on_checkpoint_restore({})
+ assert mgr._session is original_session
+ assert mgr._session.session_id == original_session_id
+
+
# endregion
diff --git a/python/pyproject.toml b/python/pyproject.toml
index af80756bed..6bd15774a9 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -184,7 +184,6 @@ omit = [
[tool.pyright]
include = ["agent_framework*"]
-exclude = ["**/tests/**", "**/.venv/**", "packages/devui/frontend/**"]
typeCheckingMode = "strict"
reportUnnecessaryIsInstance = false
reportMissingTypeStubs = false
diff --git a/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py b/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py
index f31e01ea1c..f7662d1e2f 100644
--- a/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py
+++ b/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py
@@ -61,7 +61,7 @@ async def main() -> None:
print(f"Creating memory store '{memory_store_name}'...")
try:
# Create a memory store
- memory_store = await project_client.memory_stores.create(
+ memory_store = await project_client.beta.memory_stores.create(
name=memory_store_name,
description="Memory store for Agent Framework with FoundryMemoryProvider",
definition=memory_store_definition,
@@ -126,7 +126,7 @@ async def main() -> None:
print(f"Agent: {result3}\n")
print(f"Stored memories from: {memory_store.name} ({memory_store.id})")
- res = await project_client.memory_stores.search_memories(name=memory_store.name, scope="user_123")
+ res = await project_client.beta.memory_stores.search_memories(name=memory_store.name, scope="user_123")
for memory in res.memories:
print(f"Memory: {memory.memory_item.content}")
@@ -134,7 +134,7 @@ async def main() -> None:
print(f"An error occurred: {e}")
finally:
- await project_client.memory_stores.delete(memory_store_name)
+ await project_client.beta.memory_stores.delete(memory_store_name)
print("==========================================")
print("Memory store deleted")
diff --git a/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py b/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py
index e08dfcc1bc..9efa5592c7 100644
--- a/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py
+++ b/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py
@@ -8,7 +8,7 @@ from typing import Annotated
from agent_framework import tool
from agent_framework.azure import AzureAIProjectAgentProvider
from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import AgentReference, PromptAgentDefinition
+from azure.ai.projects.models import PromptAgentDefinition
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
from pydantic import Field
@@ -116,7 +116,7 @@ async def get_agent_by_name_example() -> None:
async def get_agent_by_reference_example() -> None:
"""Example of using provider.get_agent(reference=...) to retrieve a specific agent version.
- This method fetches a specific version of an agent using an AgentReference.
+ This method fetches a specific version of an agent using a reference mapping.
Use this when you need to use a particular version of an agent.
"""
print("=== provider.get_agent(reference=...) Example ===")
@@ -136,9 +136,9 @@ async def get_agent_by_reference_example() -> None:
)
try:
- # Get the agent using an AgentReference with specific version
+ # Get the agent using a reference mapping with specific version
provider = AzureAIProjectAgentProvider(project_client=project_client)
- reference = AgentReference(name=created_agent.name, version=created_agent.version)
+ reference = {"name": created_agent.name, "version": created_agent.version}
agent = await provider.get_agent(reference=reference)
print(f"Retrieved agent: {agent.name} (version via reference)")
diff --git a/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py b/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py
index 2d1cb43c30..9377a78214 100644
--- a/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py
+++ b/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py
@@ -43,7 +43,7 @@ async def main() -> None:
options=MemoryStoreDefaultOptions(user_profile_enabled=True, chat_summary_enabled=True),
)
- memory_store = await project_client.memory_stores.create(
+ memory_store = await project_client.beta.memory_stores.create(
name=memory_store_name,
description="Memory store for Agent Framework conversations",
definition=memory_store_definition,
@@ -57,7 +57,7 @@ async def main() -> None:
instructions="""You are a helpful assistant that remembers past conversations.
Use the memory search tool to recall relevant information from previous interactions.""",
tools={
- "type": "memory_search",
+ "type": "memory_search_preview",
"memory_store_name": memory_store.name,
"scope": "user_123",
"update_delay": 1, # Wait 1 second before updating memories (use higher value in production)
@@ -84,7 +84,7 @@ async def main() -> None:
# Clean up - delete the memory store
async with AIProjectClient(endpoint=endpoint, credential=credential) as project_client:
- await project_client.memory_stores.delete(memory_store_name)
+ await project_client.beta.memory_stores.delete(memory_store_name)
print("Memory store deleted")
diff --git a/python/uv.lock b/python/uv.lock
index 15b3f18c46..415aa04f2b 100644
--- a/python/uv.lock
+++ b/python/uv.lock
@@ -401,7 +401,7 @@ requires-dist = [
{ name = "agent-framework-orchestrations", marker = "extra == 'all'", editable = "packages/orchestrations" },
{ name = "agent-framework-purview", marker = "extra == 'all'", editable = "packages/purview" },
{ name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" },
- { name = "azure-ai-projects", specifier = "==2.0.0b3" },
+ { name = "azure-ai-projects", specifier = "==2.0.0b4" },
{ name = "azure-identity", specifier = ">=1,<2" },
{ name = "mcp", extras = ["ws"], specifier = ">=1.24.0,<2" },
{ name = "openai", specifier = ">=1.99.0" },
@@ -1014,7 +1014,7 @@ wheels = [
[[package]]
name = "azure-ai-projects"
-version = "2.0.0b3"
+version = "2.0.0b4"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "azure-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -1022,10 +1022,11 @@ dependencies = [
{ name = "azure-storage-blob", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "isodate", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
{ name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/24/e0/3512d3f07e9dd2eb4af684387c31598c435bd87833b6a81850972963cb9c/azure_ai_projects-2.0.0b3.tar.gz", hash = "sha256:6d09ad110086e450a47b991ee8a3644f1be97fa3085d5981d543f900d78f4505", size = 431749, upload-time = "2026-01-06T05:31:25.849Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/24/e9/1cb8e95a19fbf174cfd7b30368a011b3e17503928b7801b8d9129b7cc59b/azure_ai_projects-2.0.0b4.tar.gz", hash = "sha256:b6082eacf0a11db59ad4c48cb7962f5204b9a0391000bc22421236f229ff783a", size = 477764, upload-time = "2026-02-24T17:57:52.489Z" }
wheels = [
- { url = "https://files.pythonhosted.org/packages/4e/b6/8fbd4786bb5c0dd19eaff86ddce0fbfb53a6f90d712038272161067a076a/azure_ai_projects-2.0.0b3-py3-none-any.whl", hash = "sha256:3b3048a3ba3904d556ba392b7bd20b6e84c93bb39df6d43a6470cdb0ad08af8c", size = 240717, upload-time = "2026-01-06T05:31:27.716Z" },
+ { url = "https://files.pythonhosted.org/packages/27/6e/6445d510a8cb6a54f57e4344c14d825c37c5146fa69ccf9d9d15a29d23e2/azure_ai_projects-2.0.0b4-py3-none-any.whl", hash = "sha256:f4cf1615bd815744ddce304b97eea9456b7f6f0bd8725547c4e54e3a67534635", size = 231920, upload-time = "2026-02-24T17:57:53.917Z" },
]
[[package]]
@@ -1408,7 +1409,7 @@ name = "clr-loader"
version = "0.2.10"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
]
sdist = { url = "https://files.pythonhosted.org/packages/18/24/c12faf3f61614b3131b5c98d3bf0d376b49c7feaa73edca559aeb2aee080/clr_loader-0.2.10.tar.gz", hash = "sha256:81f114afbc5005bafc5efe5af1341d400e22137e275b042a8979f3feb9fc9446", size = 83605, upload-time = "2026-01-03T23:13:06.984Z" }
wheels = [
@@ -1887,7 +1888,7 @@ name = "exceptiongroup"
version = "1.3.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" },
+ { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" },
]
sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
wheels = [
@@ -4654,8 +4655,8 @@ name = "powerfx"
version = "0.0.34"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
- { name = "pythonnet", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
+ { name = "pythonnet", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
]
sdist = { url = "https://files.pythonhosted.org/packages/9f/fb/6c4bf87e0c74ca1c563921ce89ca1c5785b7576bca932f7255cdf81082a7/powerfx-0.0.34.tar.gz", hash = "sha256:956992e7afd272657ed16d80f4cad24ec95d9e4a79fb9dfa4a068a09e136af32", size = 3237555, upload-time = "2025-12-22T15:50:59.682Z" }
wheels = [
@@ -5318,7 +5319,7 @@ name = "pythonnet"
version = "3.0.5"
source = { registry = "https://pypi.org/simple" }
dependencies = [
- { name = "clr-loader", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+ { name = "clr-loader", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
]
sdist = { url = "https://files.pythonhosted.org/packages/9a/d6/1afd75edd932306ae9bd2c2d961d603dc2b52fcec51b04afea464f1f6646/pythonnet-3.0.5.tar.gz", hash = "sha256:48e43ca463941b3608b32b4e236db92d8d40db4c58a75ace902985f76dac21cf", size = 239212, upload-time = "2024-12-13T08:30:44.393Z" }
wheels = [