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

python/packages/core/AGENTS.md

@@ -94,11 +94,11 @@ agent_framework/
 
 ### File Access Harness (`_harness/_file_access.py`)
 
-- **`AgentFileStore`** - Abstract async store backing the file-access harness. Implementations expose `write_file`, `read_file`, `delete_file`, `list_files`, `file_exists`, `search_files`, and `create_directory` over forward-slash relative paths.
+- **`AgentFileStore`** - Abstract async store backing the file-access harness. Implementations expose `write_file`, `read_file`, `delete_file`, `list_files`, `list_directories`, `file_exists`, `search_files`, and `create_directory` over forward-slash relative paths. `list_files`/`list_directories` return only direct children; `search_files` accepts a keyword-only `recursive` flag (default `False`) and, when `recursive=True`, walks all descendants and returns `file_name` values relative to the search directory.
 - **`InMemoryAgentFileStore`** - Dict-backed store suitable for tests and lightweight scenarios.
 - **`FileSystemAgentFileStore`** - Disk-backed store rooted under a configurable directory. Enforces relative-path normalization, root containment, and rejects symlink/reparse-point segments to prevent escape.
 - **`FileSearchResult`** / **`FileSearchMatch`** - `SerializationMixin` DTOs returned by `search_files`, carrying the matching file name, a context snippet, and the matching lines with 1-based line numbers.
-- **`FileAccessProvider`** - `ContextProvider` that adds shared file-access tools (`file_access_save_file`, `file_access_read_file`, `file_access_delete_file`, `file_access_list_files`, `file_access_search_files`) plus default usage instructions to each invocation. Unlike `MemoryContextProvider`, the store is intentionally shared across sessions and agents.
+- **`FileAccessProvider`** - `ContextProvider` that adds shared file-access tools (`file_access_save_file`, `file_access_read_file`, `file_access_delete_file`, `file_access_list_files`, `file_access_list_subdirectories`, `file_access_search_files`) plus default usage instructions to each invocation. `file_access_list_files`/`file_access_list_subdirectories` enumerate direct children (files / subdirectories) so the agent can walk the tree level by level; `file_access_search_files` searches recursively from the store root and returns store-root-relative `file_name` paths, scoped via an `fnmatch` glob (where `*` crosses `/`, e.g. `*.md`, `reports/*`). Unlike `MemoryContextProvider`, the store is intentionally shared across sessions and agents.
 
 ### Tool Approval Harness (`_harness/_tool_approval.py`)
 
@@ -116,6 +116,11 @@ agent_framework/
   available, approval requests for known non-approval-required tools are treated as already approved, hidden, stored
   in session state keyed to the visible approval request ids from that batch, and reinjected only when that visible
   approval flow resumes.
+### Agent Loop (`_harness/_loop.py`)
+
+- **`AgentLoopMiddleware`** - `AgentMiddleware` that re-runs an agent in a loop by calling `call_next()` repeatedly (the pipeline re-reads `context.messages` each time). One configurable class covers two patterns: a required user `should_continue` predicate (sync or async, the first positional/keyword arg), and a chat-client judge built via the `.with_judge(...)` factory (a second chat client decides whether the original request was answered; loops while it is *not*, using a `JudgeVerdict` structured-output response — internally just an async `should_continue` predicate). The constructor covers the predicate pattern directly; only the judge has a convenience classmethod factory (`.with_judge(judge_client, ...)`) that forwards to `__init__`. Supports both streaming and non-streaming runs. By default a non-streaming run returns an aggregated `AgentResponse` containing every iteration's messages plus the injected `next_message` "nudge" messages (as `user` messages); set `return_final_only=True` to return only the last iteration's response. Streaming runs always yield each iteration's updates and emit the injected nudge messages as `user` updates between iterations (the `return_final_only` flag has no effect on streaming, and the final response reflects the last iteration; `MiddlewareTermination` is handled cleanly). `should_continue` is required; other constructor args are optional: `max_iterations` (safety cap; defaults to `DEFAULT_MAX_ITERATIONS`=10, explicit `None`→unbounded, positive int caps; `.with_judge` uses `DEFAULT_JUDGE_MAX_ITERATIONS`=5 as its default), `next_message` (defaults to a short "continue" nudge), `return_final_only`, and `additional_instructions` (an extra `system` message injected ahead of the input before the agent runs — becomes part of the original messages so it survives `fresh_context` resets and persists via a session). The judge is configured only through `.with_judge` (`judge_client`/`instructions`/`criteria`), not the constructor, and its `reasoning` is fed back to the agent as the next iteration's input; the judge forwards the original request messages and the agent's latest response messages verbatim so multi-modal content is preserved. `criteria` (a `list[str]`) is both injected as the agent's `additional_instructions` and rendered into the judge instructions wherever the `{{criteria}}` placeholder (`CRITERIA_PLACEHOLDER`) appears (`DEFAULT_JUDGE_INSTRUCTIONS` ends with it; custom `instructions` may include it, and it is stripped when no criteria are given). The `should_continue`/`next_message` callables are invoked with keyword args (`iteration`, `last_result`, `messages`, `original_messages`, `session`, `agent`, `progress`, `feedback`) and may be sync or async; declare only what you need plus `**kwargs`. `should_continue` may return a plain `bool` or a `(bool, str | None)` tuple whose second item is feedback surfaced to `next_message`/`record_feedback` via the `feedback` kwarg (the judge uses this to relay its `reasoning`). Stop precedence per iteration is `max_iterations` → `should_continue`, evaluated before `record_feedback` so the feedback is available to it.
+  - **Feedback tracking** - `record_feedback` captures a per-iteration progress entry (called with the loop kwargs; if it returns a truthy string the entry is appended, otherwise the agent's response text is used as the fallback entry). The accumulated log is exposed to every callback via the `progress` keyword (a per-iteration copy of prior entries) and, when `inject_progress=True` (default), injected into the next iteration's input as a `user` message (the full log without a session, only the latest entry with a session to avoid duplicating history). `fresh_context=True` restarts each iteration from the original task plus the progress log; when a session is attached it is snapshotted (`to_dict()`) before the loop and restored (`from_dict` + field copy) between iterations so the local transcript and any service-side conversation id reset too (in-loop working-state is discarded, pre-loop state preserved, continuity carried only by the progress log).
+- **`todos_remaining(provider)`** / **`background_tasks_running(provider)`** - Helper factories returning `should_continue` predicates that loop while a `TodoProvider` has open items, or while a `BackgroundAgentsProvider`'s persisted state shows running tasks.
 
 ### Workflows (`_workflows/`)
 

python/packages/core/agent_framework/__init__.py

@@ -102,6 +102,12 @@ from ._harness._file_access import (
     FileSystemAgentFileStore,
     InMemoryAgentFileStore,
 )
+from ._harness._loop import (
+    AgentLoopMiddleware,
+    JudgeVerdict,
+    background_tasks_running,
+    todos_remaining,
+)
 from ._harness._memory import (
     DEFAULT_MEMORY_SOURCE_ID,
     MemoryContextProvider,
@@ -363,6 +369,7 @@ __all__ = [
     "AgentExecutorResponse",
     "AgentFileStore",
     "AgentFrameworkException",
+    "AgentLoopMiddleware",
     "AgentMiddleware",
     "AgentMiddlewareLayer",
     "AgentMiddlewareTypes",
@@ -455,6 +462,7 @@ __all__ = [
     "InlineSkill",
     "InlineSkillResource",
     "InlineSkillScript",
+    "JudgeVerdict",
     "LocalEvaluator",
     "MCPSkill",
     "MCPSkillResource",
@@ -558,6 +566,7 @@ __all__ = [
     "agent_middleware",
     "annotate_message_groups",
     "apply_compaction",
+    "background_tasks_running",
     "chat_middleware",
     "create_always_approve_tool_response",
     "create_always_approve_tool_with_arguments_response",
@@ -588,6 +597,7 @@ __all__ = [
     "response_handler",
     "set_agent_mode",
     "step",
+    "todos_remaining",
     "tool",
     "tool_call_args_match",
     "tool_called_check",

python/packages/core/agent_framework/_harness/_file_access.py

@@ -5,9 +5,10 @@
 Unlike :class:`~agent_framework.MemoryContextProvider`, which provides
 session-scoped memory that may be isolated per session, :class:`FileAccessProvider`
 operates on a shared, persistent storage area whose contents are visible across
-sessions and agents. The provider exposes five tools — ``file_access_save_file``,
+sessions and agents. The provider exposes six tools — ``file_access_save_file``,
 ``file_access_read_file``, ``file_access_delete_file``, ``file_access_list_files``,
-and ``file_access_search_files`` — by registering them on the per-invocation
+``file_access_list_subdirectories``, and ``file_access_search_files`` — by
+registering them on the per-invocation
 :class:`~agent_framework.SessionContext` in :meth:`FileAccessProvider.before_run`.
 
 The store abstraction is generic so callers can plug in in-memory, local-disk, or
@@ -48,7 +49,11 @@ DEFAULT_FILE_ACCESS_INSTRUCTIONS = (
     "Use these tools to read input data provided by the user, write output "
     "artifacts, and manage any files the user has asked you to work with.\n\n"
     "- Never delete or overwrite existing files unless the user has explicitly "
-    "asked you to do so."
+    "asked you to do so.\n"
+    "- Files may be organized into subdirectories. Use `file_access_list_files` "
+    "and `file_access_list_subdirectories` to explore the tree level by level, "
+    "or `file_access_search_files` to search file contents recursively across "
+    "the whole store."
 )
 
 # Maximum number of characters of context to include on either side of the first
@@ -178,10 +183,16 @@ def _normalize_relative_path(path: str, *, is_directory: bool = False) -> str:
 def _matches_glob(file_name: str, pattern: str | None) -> bool:
     """Return whether ``file_name`` matches the optional glob pattern (case-insensitive).
 
-    When ``pattern`` is ``None`` or blank this returns True so callers can skip
-    filtering by passing nothing. Matching uses :func:`fnmatch.fnmatchcase` over a
-    lowercased pattern/name pair to give consistent results across operating
-    systems (``fnmatch.fnmatch`` is case-sensitive on POSIX but not on Windows).
+    ``file_name`` is the forward-slash path of a file relative to the search
+    directory (for a direct child this is just its basename; for a recursive
+    search it may contain ``/`` separators). When ``pattern`` is ``None`` or blank
+    this returns True so callers can skip filtering by passing nothing. Matching
+    uses :func:`fnmatch.fnmatchcase` over a lowercased pattern/name pair to give
+    consistent results across operating systems (``fnmatch.fnmatch`` is
+    case-sensitive on POSIX but not on Windows). Note that with ``fnmatch`` a
+    ``*`` matches any characters **including** ``/``, so ``"*.md"`` matches
+    markdown files at any depth and ``"reports/*"`` matches everything under
+    ``reports``.
     """
     if pattern is None or not pattern.strip():
         return True
@@ -418,6 +429,18 @@ class AgentFileStore(ABC):
             The list of file names (not full paths) in the specified directory.
         """
 
+    @abstractmethod
+    async def list_directories(self, directory: str = "") -> list[str]:
+        """List the direct child subdirectory names of ``directory``.
+
+        Args:
+            directory: The relative directory path to list. Use ``""`` for the root.
+
+        Returns:
+            The list of subdirectory names (not full paths) directly contained in
+            the specified directory.
+        """
+
     @abstractmethod
     async def file_exists(self, path: str) -> bool:
         """Return whether a file exists at ``path``.
@@ -432,6 +455,8 @@ class AgentFileStore(ABC):
         directory: str,
         regex_pattern: str,
         file_pattern: str | None = None,
+        *,
+        recursive: bool = False,
     ) -> list[FileSearchResult]:
         """Search files in ``directory`` for content matching ``regex_pattern``.
 
@@ -441,12 +466,19 @@ class AgentFileStore(ABC):
                 (case-insensitive). For example, ``"error|warning"`` matches lines
                 containing ``"error"`` or ``"warning"``.
             file_pattern: An optional glob pattern (case-insensitive) used to
-                filter which files are searched. When ``None`` or blank, every
-                file in the directory is searched.
+                filter which files are searched. The pattern is matched against
+                each file's path **relative to** ``directory`` (forward slashes).
+                When ``None`` or blank, every file in scope is searched.
+
+        Keyword Args:
+            recursive: When ``False`` (default) only the direct children of
+                ``directory`` are searched. When ``True`` every descendant file is
+                searched.
 
         Returns:
             The list of files whose content matched, with snippet and matching
-            line metadata.
+            line metadata. Each result's ``file_name`` is the path relative to
+            ``directory`` (forward slashes).
         """
 
     @abstractmethod
@@ -530,6 +562,38 @@ class InMemoryAgentFileStore(AgentFileStore):
             results.append(display[len(prefix) :])
         return results
 
+    async def list_directories(self, directory: str = "") -> list[str]:
+        """Return the direct child subdirectory names of ``directory``.
+
+        A subdirectory is the first path segment of any stored key whose
+        remainder (after the directory prefix) still contains a ``/`` separator.
+        Distinct first segments are collected, preserving the *original-case*
+        display name and de-duplicating case-insensitively, mirroring the
+        case-preserving behaviour of :class:`FileSystemAgentFileStore`.
+        """
+        prefix = _normalize_relative_path(directory, is_directory=True).lower()
+        if prefix and not prefix.endswith("/"):
+            prefix += "/"
+        async with self._lock:
+            entries = [(key, display) for key, (display, _) in self._files.items()]
+        results: list[str] = []
+        seen: set[str] = set()
+        for key, display in entries:
+            if not key.startswith(prefix):
+                continue
+            remainder = key[len(prefix) :]
+            separator_index = remainder.find("/")
+            if separator_index <= 0:
+                continue
+            segment_key = remainder[:separator_index]
+            if segment_key in seen:
+                continue
+            seen.add(segment_key)
+            # ``display`` is the original-case normalized path; take the matching
+            # first segment after the (case-insensitive) prefix.
+            results.append(display[len(prefix) : len(prefix) + separator_index])
+        return results
+
     async def file_exists(self, path: str) -> bool:
         """Return whether the file exists."""
         key = self._key(path)
@@ -541,6 +605,8 @@ class InMemoryAgentFileStore(AgentFileStore):
         directory: str,
         regex_pattern: str,
         file_pattern: str | None = None,
+        *,
+        recursive: bool = False,
     ) -> list[FileSearchResult]:
         """Search file contents for ``regex_pattern`` matches.
 
@@ -548,7 +614,10 @@ class InMemoryAgentFileStore(AgentFileStore):
         to a worker thread with a bounded timeout so a pathological pattern
         cannot stall the event loop. Returned :class:`FileSearchResult`
         instances use the *original-case* file names so the result mirrors
-        what :class:`FileSystemAgentFileStore` would produce.
+        what :class:`FileSystemAgentFileStore` would produce. The glob and each
+        result's ``file_name`` are relative to ``directory``; when ``recursive``
+        is ``True`` all descendants are searched and the relative path may
+        contain ``/`` separators.
         """
         prefix = _normalize_relative_path(directory, is_directory=True).lower()
         if prefix and not prefix.endswith("/"):
@@ -564,7 +633,7 @@ class InMemoryAgentFileStore(AgentFileStore):
                 if not key.startswith(prefix):
                     continue
                 relative_key = key[len(prefix) :]
-                if "/" in relative_key:
+                if not recursive and "/" in relative_key:
                     continue
                 relative_display = display[len(prefix) :]
                 if not _matches_glob(relative_display, file_pattern):
@@ -795,6 +864,28 @@ class FileSystemAgentFileStore(AgentFileStore):
                 names.append(entry.name)
         return names
 
+    async def list_directories(self, directory: str = "") -> list[str]:
+        """Return the direct child subdirectory names of ``directory``.
+
+        Symlinked directories (and reparse points on Windows) are excluded so a
+        listing cannot surface a path that escapes the root. An empty list is
+        returned for a non-existent directory.
+        """
+        full_dir = self._resolve_safe_directory_path(directory)
+        return await asyncio.to_thread(self._list_directories_sync, full_dir)
+
+    @staticmethod
+    def _list_directories_sync(full_dir: Path) -> list[str]:
+        if not full_dir.is_dir():
+            return []
+        names: list[str] = []
+        for entry in full_dir.iterdir():
+            if entry.is_symlink():
+                continue
+            if entry.is_dir():
+                names.append(entry.name)
+        return names
+
     async def file_exists(self, path: str) -> bool:
         """Return whether the file exists."""
         full_path = self._resolve_safe_path(path)
@@ -809,6 +900,8 @@ class FileSystemAgentFileStore(AgentFileStore):
         directory: str,
         regex_pattern: str,
         file_pattern: str | None = None,
+        *,
+        recursive: bool = False,
     ) -> list[FileSearchResult]:
         """Search file contents for ``regex_pattern`` matches.
 
@@ -816,23 +909,50 @@ class FileSystemAgentFileStore(AgentFileStore):
         file does not abort the whole directory search). Each skip is logged at
         ``WARNING`` level and a summary is logged at ``INFO`` so operators can
         tell the difference between "no matches" and "the corpus was largely
-        not searchable".
+        not searchable". The glob and each result's ``file_name`` are the file's
+        path relative to ``directory`` (forward slashes); when ``recursive`` is
+        ``True`` all descendant files are searched, otherwise only the direct
+        children.
         """
         full_dir = self._resolve_safe_directory_path(directory)
         regex = _compile_search_regex(regex_pattern)
-        return await _run_search_with_timeout(lambda: self._search_files_sync(full_dir, regex, file_pattern))
+        return await _run_search_with_timeout(lambda: self._search_files_sync(full_dir, regex, file_pattern, recursive))
 
     @staticmethod
-    def _search_files_sync(full_dir: Path, regex: re.Pattern[str], file_pattern: str | None) -> list[FileSearchResult]:
+    def _enumerate_search_files(full_dir: Path, recursive: bool) -> list[tuple[str, Path]]:
+        """Enumerate ``(relative_name, path)`` for files to search under ``full_dir``.
+
+        Symlinked files and symlinked directories (reparse points on Windows)
+        are skipped so the search cannot read or descend outside the root.
+        ``relative_name`` is the file's path relative to ``full_dir`` using
+        forward slashes.
+        """
+        found: list[tuple[str, Path]] = []
+        directories: list[Path] = [full_dir]
+        while directories:
+            current = directories.pop()
+            for entry in current.iterdir():
+                if entry.is_symlink():
+                    continue
+                if entry.is_dir():
+                    if recursive:
+                        directories.append(entry)
+                    continue
+                if entry.is_file():
+                    relative_name = entry.relative_to(full_dir).as_posix()
+                    found.append((relative_name, entry))
+        return found
+
+    @staticmethod
+    def _search_files_sync(
+        full_dir: Path, regex: re.Pattern[str], file_pattern: str | None, recursive: bool
+    ) -> list[FileSearchResult]:
         if not full_dir.is_dir():
             return []
         results: list[FileSearchResult] = []
         skipped: list[str] = []
-        for entry in full_dir.iterdir():
-            if entry.is_symlink() or not entry.is_file():
-                continue
-            file_name = entry.name
-            if not _matches_glob(file_name, file_pattern):
+        for relative_name, entry in FileSystemAgentFileStore._enumerate_search_files(full_dir, recursive):
+            if not _matches_glob(relative_name, file_pattern):
                 continue
             try:
                 file_content = entry.read_text(encoding="utf-8")
@@ -841,9 +961,9 @@ class FileSystemAgentFileStore(AgentFileStore):
                 # un-decodable entry doesn't abort the whole directory search.
                 # Log per file so operators can audit which files were skipped.
                 logger.warning("Skipping non-UTF-8 file during search: %s", entry)
-                skipped.append(file_name)
+                skipped.append(relative_name)
                 continue
-            result = _search_file_content(file_name, file_content, regex)
+            result = _search_file_content(relative_name, file_content, regex)
             if result is not None:
                 results.append(result)
         if skipped:
@@ -865,15 +985,18 @@ class FileSystemAgentFileStore(AgentFileStore):
 class FileAccessProvider(ContextProvider):
     """Context provider that gives an agent CRUD/search access to a shared file store.
 
-    The provider exposes five tools to the agent via the per-invocation
+    The provider exposes six tools to the agent via the per-invocation
     :class:`~agent_framework.SessionContext`:
 
     - ``file_access_save_file`` — Save a file (refuses to overwrite by default).
     - ``file_access_read_file`` — Read the content of a file by name.
     - ``file_access_delete_file`` — Delete a file by name.
-    - ``file_access_list_files`` — List all file names at the store root.
-    - ``file_access_search_files`` — Search file contents using a case-insensitive
-      regex, optionally filtered by a glob pattern over file names.
+    - ``file_access_list_files`` — List the direct child file names of a directory.
+    - ``file_access_list_subdirectories`` — List the direct child subdirectory
+      names of a directory.
+    - ``file_access_search_files`` — Recursively search file contents from the
+      store root using a case-insensitive regex, optionally filtered by a glob
+      pattern over the store-root-relative file paths.
 
     Unlike :class:`~agent_framework.MemoryContextProvider`, which provides
     session-scoped memory that may be isolated per session,
@@ -976,17 +1099,45 @@ class FileAccessProvider(ContextProvider):
             except OSError as exc:
                 return f"Could not list directory '{directory or ''}': {exc.strerror or exc}"
 
+        @tool(name="file_access_list_subdirectories", approval_mode="never_require")
+        async def file_access_list_subdirectories(directory: str | None = None) -> list[str] | str:
+            """List the direct child subdirectory names of a directory.
+
+            Omit ``directory`` (or pass an empty string) to list the root.
+            To enumerate subdirectories of a subdirectory, pass its relative path, for example
+            ``"reports"`` or ``"reports/2024"``.
+            Use this together with file_access_list_files to explore the directory tree level by level.
+            """
+            target = directory if directory and directory.strip() else ""
+            try:
+                return await self.store.list_directories(target)
+            except ValueError as exc:
+                return f"Could not list directory '{directory or ''}': {exc}"
+            except OSError as exc:
+                return f"Could not list directory '{directory or ''}': {exc.strerror or exc}"
+
         @tool(name="file_access_search_files", approval_mode="never_require")
         async def file_access_search_files(
             regex_pattern: str,
             file_pattern: str | None = None,
-            directory: str | None = None,
         ) -> list[dict[str, Any]] | str:
-            """Search file contents using a regular expression pattern (case-insensitive). Optionally filter which files to search using a glob pattern (e.g., "*.md", "research*"). Optionally scope the search to a subdirectory by passing its relative path; omit ``directory`` (or pass an empty string) to search the root. Returns matching file names, snippets, and matching lines with line numbers. The regex_pattern must be 256 characters or fewer."""  # noqa: E501
+            """Search the contents of all files in the store using a case-insensitive regular expression.
+
+            The search runs recursively across all subdirectories.
+            Optionally filter which files to search using a glob pattern matched against each file's
+            path relative to the store root.
+            The glob uses fnmatch semantics where ``*`` matches any characters including ``/``: use
+            ``"*.md"`` to match markdown files at any depth,
+            or ``"reports/*"`` to restrict the search to the ``reports`` subtree.
+            Leave empty or omit to search all files.
+            Returns matching results whose file_name values are paths relative to the store root
+            (usable with file_access_read_file),
+            along with snippets and matching lines with line numbers. The regex_pattern must be
+            256 characters or fewer.
+            """
             pattern = file_pattern if file_pattern and file_pattern.strip() else None
-            target = directory if directory and directory.strip() else ""
             try:
-                results = await self.store.search_files(target, regex_pattern, pattern)
+                results = await self.store.search_files("", regex_pattern, pattern, recursive=True)
             except ValueError as exc:
                 return f"Could not search files: {exc}"
             except OSError as exc:
@@ -1001,6 +1152,7 @@ class FileAccessProvider(ContextProvider):
                 file_access_read_file,
                 file_access_delete_file,
                 file_access_list_files,
+                file_access_list_subdirectories,
                 file_access_search_files,
             ],
         )

python/packages/core/agent_framework/_harness/_loop.py

@@ -0,0 +1,796 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""AgentLoopMiddleware: re-run an agent in a loop until a criterion is met.
+
+This module provides :class:`AgentLoopMiddleware`, an :class:`~agent_framework.AgentMiddleware`
+that repeatedly re-invokes the wrapped agent while a ``should_continue`` predicate says to keep
+going. It serves two common patterns through a single configurable class:
+
+1. A user-supplied ``should_continue`` predicate - for example, keep looping while a response does
+   not yet contain a completion marker, while a :class:`~agent_framework.TodoProvider` still has
+   open items, or while a :class:`~agent_framework.BackgroundAgentsProvider` still has running
+   tasks (see the :func:`todos_remaining` and :func:`background_tasks_running` helpers). The loop
+   can track a **feedback log** across iterations (``record_feedback``): each pass contributes an
+   entry that is exposed to every callback via the ``progress`` keyword and (by default) injected
+   into the next iteration's input. Set ``fresh_context=True`` to restart each pass from the
+   original task plus the progress log (with a session attached, the session is also snapshotted
+   before the loop and restored between iterations so no accumulated history leaks back in).
+   ``max_iterations`` bounds the loop as a safety cap.
+2. A chat-client judge (via :meth:`AgentLoopMiddleware.with_judge`) - a second chat client decides
+   whether the user's original request has been answered (via a :class:`JudgeVerdict` structured
+   output); the loop continues while the answer is "no". This is a convenience wrapper that builds an
+   async ``should_continue`` predicate, so it is a special case of (1).
+
+In every case, the input for the next iteration is controlled by the ``next_message`` callable.
+"""
+
+from __future__ import annotations
+
+import inspect
+from collections.abc import Awaitable, Callable, Sequence
+from typing import TYPE_CHECKING, Any, TypeAlias
+
+from pydantic import BaseModel, Field
+from typing_extensions import Self
+
+from .._feature_stage import ExperimentalFeature, experimental
+from .._middleware import AgentContext, AgentMiddleware, MiddlewareTermination
+from .._types import (
+    AgentResponse,
+    AgentResponseUpdate,
+    AgentRunInputs,
+    Message,
+    ResponseStream,
+    UsageDetails,
+    add_usage_details,
+    normalize_messages,
+)
+
+if TYPE_CHECKING:
+    from .._clients import SupportsChatGetResponse
+
+__all__ = [
+    "AgentLoopMiddleware",
+    "JudgeVerdict",
+    "background_tasks_running",
+    "todos_remaining",
+]
+
+DEFAULT_NEXT_MESSAGE = "Continue working on the task. If it is complete, say so."
+
+# Placeholder substituted with the rendered ``criteria`` block in judge instructions (see
+# :meth:`AgentLoopMiddleware.with_judge`). User-supplied instructions may include it to control
+# where the criteria are inserted; if absent, the criteria are not added to the judge instructions.
+CRITERIA_PLACEHOLDER = "{{criteria}}"
+
+# Verdict markers the judge is asked to emit for clients that do not honor structured output. They
+# are deliberately non-overlapping: neither marker is a substring of the other, nor of the JSON
+# field name ``answered``, so the text fallback in :func:`_build_judge_condition` cannot misclassify
+# a negative verdict (e.g. ``{"answered": false}``) as a positive one.
+JUDGE_VERDICT_DONE = "VERDICT: DONE"
+JUDGE_VERDICT_MORE = "VERDICT: MORE"
+
+DEFAULT_JUDGE_INSTRUCTIONS = (
+    "You are an evaluator. You are given a user's original request and an agent's latest response. "
+    "Decide whether the agent has fully addressed the original request. "
+    "Set 'answered' to true if the request has been fully addressed, or false if more work is still "
+    "required, and use 'reasoning' to briefly justify your decision. "
+    f"If you cannot return structured output, end your reply with a line reading exactly "
+    f"'{JUDGE_VERDICT_DONE}' when the request has been fully addressed or '{JUDGE_VERDICT_MORE}' "
+    f"when more work is still required."
+    "{{criteria}}"
+)
+
+
+def _render_criteria_block(criteria: Sequence[str] | None) -> str:
+    """Render a list of criteria into a bullet block for the judge instructions (``""`` if none)."""
+    if not criteria:
+        return ""
+    bullets = "\n".join(f"- {item}" for item in criteria)
+    return f"\n\nThe response must satisfy all of the following criteria:\n{bullets}"
+
+
+def _criteria_agent_instruction(criteria: Sequence[str]) -> str:
+    """Render the criteria into an extra instruction injected for the agent before each run."""
+    bullets = "\n".join(f"- {item}" for item in criteria)
+    return f"Your response must satisfy all of the following criteria:\n{bullets}"
+
+
+class JudgeVerdict(BaseModel):
+    """Structured verdict returned by the judge chat client."""
+
+    answered: bool = Field(
+        description=(
+            "True if the agent has fully addressed the original request and it adheres to the other "
+            "judging standards, otherwise False."
+        ),
+    )
+    reasoning: str = Field(
+        default="",
+        description="Brief justification for the verdict.",
+    )
+
+
+# Default iteration cap applied when ``max_iterations`` is not provided. Loops are bounded by
+# default to guard against runaway re-invocation; pass ``max_iterations=None`` explicitly to opt
+# into an unbounded loop.
+DEFAULT_MAX_ITERATIONS = 10
+
+# Default iteration cap for judge-driven loops. LLM-judged loops are costly and probabilistic, so
+# they are bounded by a smaller default. Pass ``max_iterations=None`` explicitly to opt into an
+# unbounded judge loop.
+DEFAULT_JUDGE_MAX_ITERATIONS = 5
+
+
+# A callable invoked between iterations. It always receives the loop keyword arguments
+# (``iteration``, ``last_result``, ``messages``, ``original_messages``, ``session``, ``agent``,
+# ``progress``, ``feedback``). Callers declare only the keywords they need plus ``**kwargs`` to
+# ignore the rest. ``should_continue`` may return a plain ``bool`` (continue/stop) or a
+# ``(bool, str | None)`` tuple whose second item is feedback surfaced to the ``next_message`` and
+# ``record_feedback`` callables via the ``feedback`` keyword argument.
+ShouldContinueResult: TypeAlias = "bool | tuple[bool, str | None]"
+ShouldContinueCallable = Callable[..., "ShouldContinueResult | Awaitable[ShouldContinueResult]"]
+NextMessageCallable = Callable[..., "AgentRunInputs | Awaitable[AgentRunInputs | None] | None"]
+
+# A callable invoked once per work iteration to capture a progress-log entry from that iteration. It
+# receives the loop keyword arguments and returns a string entry (appended to the log) or ``None``
+# (record nothing for that iteration).
+FeedbackCallable = Callable[..., "str | Awaitable[str | None] | None"]
+
+
+async def _maybe_await(value: Any) -> Any:
+    """Await ``value`` if it is awaitable, otherwise return it as-is."""
+    if inspect.isawaitable(value):
+        return await value
+    return value
+
+
+def _build_judge_condition(
+    judge_client: SupportsChatGetResponse,
+    instructions: str,
+) -> tuple[ShouldContinueCallable, NextMessageCallable]:
+    """Build the ``should_continue`` predicate and ``next_message`` callable for a judge loop.
+
+    The judge is called directly (no agent tools, session, or middleware) with fresh messages, so
+    the loop's evaluation cannot recurse back through the agent pipeline. The original input messages
+    are forwarded verbatim (rather than collapsed to text) so multi-modal requests are preserved. The
+    judge is asked for a :class:`JudgeVerdict` structured output; if the client does not honor
+    structured output the verdict falls back to the explicit, non-overlapping ``VERDICT: DONE`` /
+    ``VERDICT: MORE`` markers (``MORE`` wins, keeping the loop running, when the marker is ambiguous
+    or absent).
+
+    The predicate returns a ``(continue, reasoning)`` tuple; the loop surfaces that ``reasoning`` to
+    the next-message callable as the ``feedback`` keyword argument, which feeds it back to the agent
+    so it knows *why* its previous answer was judged incomplete.
+    """
+
+    async def _judge(
+        *, last_result: AgentResponse, original_messages: list[Message], **kwargs: Any
+    ) -> tuple[bool, str | None]:
+        judge_messages = [
+            Message(role="system", contents=[instructions]),
+            Message(
+                role="user",
+                contents=["Evaluate the agent's work. The user's original request follows:"],
+            ),
+            *original_messages,
+            Message(role="user", contents=["The agent's latest response was:"]),
+            *last_result.messages,
+            Message(role="user", contents=["Has the original request been fully addressed?"]),
+        ]
+        response = await judge_client.get_response(judge_messages, options={"response_format": JudgeVerdict})
+        verdict = response.value
+        if isinstance(verdict, JudgeVerdict):
+            answered = verdict.answered
+            reasoning = verdict.reasoning
+        else:
+            # Fallback for clients that do not honor structured output: look for the explicit,
+            # non-overlapping verdict markers. ``FAIL`` (more work needed) takes precedence so an
+            # ambiguous or marker-less reply keeps looping rather than stopping on an incomplete
+            # answer.
+            text = response.text.upper()
+            # ``MORE`` (more work needed) takes precedence so an ambiguous reply keeps looping.
+            answered = False if JUDGE_VERDICT_MORE in text else JUDGE_VERDICT_DONE in text
+            reasoning = response.text.strip()
+        # Continue looping while the request is not yet answered, surfacing the reasoning as feedback.
+        return (not answered), (reasoning or None)
+
+    def _next_message(*, feedback: str | None = None, **kwargs: Any) -> AgentRunInputs:
+        # Feed the judge's reasoning back to the agent so the next iteration addresses the gap.
+        if feedback:
+            return (
+                "An evaluator reviewed your previous response and judged that it does not yet fully "
+                f"address the original request.\n\nEvaluator feedback: {feedback}\n\n"
+                "Revise and continue so the original request is fully addressed."
+            )
+        return DEFAULT_NEXT_MESSAGE
+
+    return _judge, _next_message
+
+
+@experimental(feature_id=ExperimentalFeature.HARNESS)
+class AgentLoopMiddleware(AgentMiddleware):
+    """Re-run an agent in a loop until a criterion is met (or never).
+
+    This middleware repeatedly invokes the wrapped agent. After each run it decides whether to run
+    again based on ``should_continue`` and ``max_iterations``, and uses ``next_message`` to build
+    the input for the next iteration. Use :meth:`with_judge` to drive the loop with a chat-client
+    judge instead of a hand-written predicate.
+
+    By default a non-streaming run returns an aggregated :class:`~agent_framework.AgentResponse`
+    containing every iteration's messages plus the injected ``next_message`` "nudge" messages (set
+    ``return_final_only=True`` to return only the last iteration's response). Streaming runs always
+    yield each iteration's updates and emit the injected nudge messages as ``user`` updates between
+    iterations.
+
+    The ``should_continue`` and ``next_message`` callables are invoked with keyword arguments, so a
+    caller only needs to declare the ones it uses plus ``**kwargs``. The keywords are:
+
+    - ``iteration`` (int): the number of completed runs so far (1-based after the first run).
+    - ``last_result`` (AgentResponse): the result of the iteration that just completed.
+    - ``messages`` (list[Message]): the messages used for the iteration that just completed.
+    - ``original_messages`` (list[Message]): the input used for the first iteration.
+    - ``session`` (AgentSession | None): the active session, used by the provider helpers.
+    - ``agent``: the agent being looped.
+    - ``progress`` (list[str]): the feedback log accumulated so far (see ``record_feedback``).
+    - ``feedback`` (str | None): the feedback string returned by ``should_continue`` for this
+      iteration (``None`` when it returned a plain bool). ``should_continue`` may return either a
+      ``bool`` or a ``(bool, str | None)`` tuple; the string is surfaced here so ``next_message``
+      and ``record_feedback`` can reference it.
+
+    Examples:
+        .. code-block:: python
+
+            from agent_framework import Agent, AgentResponse
+            from agent_framework._harness._loop import AgentLoopMiddleware
+
+
+            async def should_continue(*, iteration: int, last_result: AgentResponse, **kwargs) -> bool:
+                return iteration < 3 and "DONE" not in last_result.text
+
+
+            agent = Agent(client=client, middleware=[AgentLoopMiddleware(should_continue)])
+
+    Note:
+        ``max_iterations`` acts as a safety cap and defaults to ``DEFAULT_MAX_ITERATIONS`` (10). Pass
+        an explicit ``None`` to make the loop unbounded, in which case it relies entirely on
+        ``should_continue`` to stop, so make sure the predicate can eventually return ``False``.
+    """
+
+    def __init__(
+        self,
+        should_continue: ShouldContinueCallable,
+        *,
+        max_iterations: int | None = DEFAULT_MAX_ITERATIONS,
+        next_message: NextMessageCallable | None = None,
+        record_feedback: FeedbackCallable | None = None,
+        inject_progress: bool = True,
+        fresh_context: bool = False,
+        return_final_only: bool = False,
+        additional_instructions: str | None = None,
+    ) -> None:
+        """Initialize the agent loop middleware.
+
+        Args:
+            should_continue: Predicate that decides whether to run the agent again. May be sync or
+                async and is called with the loop keyword arguments (``iteration``, ``last_result``,
+                ``messages``, ``original_messages``, ``session``, ``agent``, ``progress``, and
+                ``feedback`` -- see the class docstring for what each one carries; declare only the
+                ones you need plus ``**kwargs``). Return ``True``/``False`` to
+                continue/stop, or a ``(bool, str | None)`` tuple to also provide feedback; the
+                feedback string is surfaced to the ``next_message`` and ``record_feedback`` callables
+                via the ``feedback`` keyword argument. To loop on a chat-client judge instead, build
+                the middleware via :meth:`with_judge`.
+
+        Keyword Args:
+            max_iterations: Maximum number of agent runs, used as a safety cap. Defaults to
+                ``DEFAULT_MAX_ITERATIONS`` (10); pass an explicit ``None`` for an unbounded loop, or
+                a positive integer to set a custom cap. (The :meth:`with_judge` factory uses
+                ``DEFAULT_JUDGE_MAX_ITERATIONS`` (5) as its default instead.)
+            next_message: Callable that produces the input for the next iteration, called with the
+                loop keyword arguments. Defaults to a short "continue" nudge. Returning ``None``
+                reuses the previous iteration's messages verbatim (in which case the progress log is
+                *not* injected; see ``inject_progress``).
+            record_feedback: Optional callable invoked once per work iteration to capture a feedback
+                entry. Called as ``record_feedback(**loop_kwargs)`` and returns a
+                string entry appended to the progress log, or ``None`` to record nothing for that
+                iteration. When not provided, the iteration's response text (``last_result.text``) is
+                recorded instead.                 The accumulated log is exposed to every callback via the
+                ``progress`` loop keyword argument. For production loops prefer a ``record_feedback``
+                that returns a terse summary rather than relying on the full response text.
+            inject_progress: When ``True`` (default), the accumulated progress log is injected into
+                the next iteration's input as a single ``user`` message ("Progress so far: ..."). To
+                avoid duplication, only the most recent entry is injected when a session is attached
+                (the session already retains earlier turns); the full log is injected when there is
+                no session or ``fresh_context`` is set. When ``False`` the log is only exposed via the
+                ``progress`` loop keyword argument and never injected automatically.
+            fresh_context: When ``True``, each iteration starts from a clean context: ``context``
+                messages are reset to the original input messages (plus the injected progress log)
+                instead of accumulating the prior conversation. When a session is attached, the
+                session is snapshotted once before the loop and restored to that pre-loop baseline
+                before each subsequent iteration, so the local transcript and any service-side
+                conversation id are reset too and the agent does not re-read the accumulated history.
+                In-loop working-state mutations are discarded; pre-loop state is preserved; continuity
+                is carried only by the progress log.
+            return_final_only: Controls what a non-streaming run returns. When ``False`` (default),
+                the returned :class:`~agent_framework.AgentResponse` aggregates every iteration: each
+                iteration's response messages plus the injected ``next_message`` "nudge" messages
+                (as ``user`` messages), so the caller sees the full back-and-forth. When ``True``,
+                only the final iteration's :class:`~agent_framework.AgentResponse` is returned. This
+                flag has no effect on streaming runs (the stream cannot know in advance which
+                iteration is last); streaming always yields each iteration's updates and injects the
+                ``next_message`` messages as ``user`` updates between iterations.
+            additional_instructions: Optional extra instruction injected as a ``system`` message
+                ahead of the input messages before the agent runs. It becomes part of the original
+                messages, so it is preserved across ``fresh_context`` resets and (with a session)
+                persists server-side across iterations. Used by :meth:`with_judge` to tell the agent
+                about the criteria its response must satisfy, but available to any loop.
+
+        Raises:
+            ValueError: If ``max_iterations`` is not ``None`` and is less than 1.
+        """
+        if max_iterations is not None and max_iterations < 1:
+            raise ValueError("max_iterations must be None or a positive integer (>= 1).")
+
+        self.max_iterations: int | None = max_iterations
+        self.should_continue: ShouldContinueCallable = should_continue
+        self.next_message = next_message
+        self.record_feedback = record_feedback
+        self.inject_progress = inject_progress
+        self.fresh_context = fresh_context
+        self.return_final_only = return_final_only
+        self.additional_instructions = additional_instructions
+
+    @classmethod
+    def with_judge(
+        cls,
+        judge_client: SupportsChatGetResponse,
+        *,
+        criteria: Sequence[str] | None = None,
+        instructions: str | None = None,
+        max_iterations: int | None = DEFAULT_JUDGE_MAX_ITERATIONS,
+        next_message: NextMessageCallable | None = None,
+        fresh_context: bool = False,
+    ) -> Self:
+        """Create a loop that continues until a judge chat client decides the request was answered.
+
+        Convenience factory for the judge pattern: ``judge_client`` is queried with a
+        :class:`JudgeVerdict` structured-output response after each iteration and the loop continues
+        while the request is *not* answered. The judge's ``reasoning`` is fed back to the agent as
+        the next iteration's input (unless a custom ``next_message`` is provided), so the agent knows
+        why its previous answer was judged incomplete. See :meth:`__init__` for the full meaning of
+        each argument.
+
+        Args:
+            judge_client: Chat client used to judge whether the original request was answered.
+
+        Keyword Args:
+            criteria: Optional list of criteria the response must satisfy. When provided, they are
+                (1) injected as an extra ``system`` instruction for the agent before it runs (via
+                ``additional_instructions``) and (2) rendered into the judge instructions wherever
+                the ``{{criteria}}`` placeholder appears (``CRITERIA_PLACEHOLDER``).
+            instructions: Optional system instructions for the judge. Defaults to
+                ``DEFAULT_JUDGE_INSTRUCTIONS``. May contain the ``{{criteria}}`` placeholder, which
+                is replaced with the rendered ``criteria`` (or removed when no criteria are given).
+            max_iterations: Maximum number of agent runs. Defaults to
+                ``DEFAULT_JUDGE_MAX_ITERATIONS`` (5); pass ``None`` for unbounded, or a positive
+                integer to set a custom cap.
+            next_message: Callable that produces the next iteration's input. Defaults to one that
+                relays the judge's ``reasoning`` back to the agent.
+            fresh_context: When ``True``, each iteration restarts from the original input messages
+                (plus the injected progress log and judge feedback) instead of accumulating the prior
+                conversation; an attached session is snapshotted before the loop and restored to that
+                baseline between iterations. See :meth:`__init__` for the full semantics. Defaults to
+                ``False``.
+        """
+        judge_instructions = (instructions or DEFAULT_JUDGE_INSTRUCTIONS).replace(
+            CRITERIA_PLACEHOLDER, _render_criteria_block(criteria)
+        )
+        should_continue, judge_next_message = _build_judge_condition(judge_client, judge_instructions)
+        return cls(
+            should_continue=should_continue,
+            max_iterations=max_iterations,
+            next_message=next_message or judge_next_message,
+            fresh_context=fresh_context,
+            additional_instructions=_criteria_agent_instruction(criteria) if criteria else None,
+        )
+
+    async def process(
+        self,
+        context: AgentContext,
+        call_next: Callable[[], Awaitable[None]],
+    ) -> None:
+        """Run the wrapped agent in a loop."""
+        if self.additional_instructions is not None:
+            # Inject the extra instruction as a system message ahead of the input so it is present
+            # on every iteration and preserved across fresh_context resets (which restart from
+            # ``original_messages``).
+            context.messages = [
+                Message(role="system", contents=[self.additional_instructions]),
+                *context.messages,
+            ]
+        original_messages = list(context.messages)
+        # For a truly fresh context per iteration the session must also be reset, otherwise the
+        # next run reloads the local transcript or re-threads the service-side conversation and the
+        # model still sees the accumulated history. Snapshot the session once here (the pre-loop
+        # baseline) and restore it before each subsequent iteration so every pass starts clean.
+        snapshot = context.session.to_dict() if self.fresh_context and context.session is not None else None
+        if context.stream:
+            self._process_streaming(context, call_next, original_messages, snapshot)
+        else:
+            await self._process_non_streaming(context, call_next, original_messages, snapshot)
+
+    @staticmethod
+    def _restore_session(session: Any, snapshot: dict[str, Any]) -> None:
+        """Restore a session in place to a previously captured ``to_dict()`` snapshot.
+
+        Re-hydrates the snapshot via :meth:`AgentSession.from_dict` and copies the mutable fields
+        (``service_session_id`` and ``state``) back onto the live ``session`` instance, so any
+        reference held by the agent/context observes the reset. ``session_id`` is preserved (the
+        snapshot carries the same id). A fresh ``from_dict`` is built on every call so repeated
+        restores from one snapshot do not alias the same state dict.
+        """
+        from .._sessions import AgentSession
+
+        restored = AgentSession.from_dict(snapshot)
+        session.service_session_id = restored.service_session_id
+        session.state = restored.state
+
+    async def _process_non_streaming(
+        self,
+        context: AgentContext,
+        call_next: Callable[[], Awaitable[None]],
+        original_messages: list[Message],
+        snapshot: dict[str, Any] | None,
+    ) -> None:
+        iteration = 0
+        work_iterations = 0
+        progress: list[str] = []
+        # Aggregated transcript across iterations: each iteration's response messages plus the
+        # injected "nudge" messages, used to build the combined response when return_final_only=False.
+        aggregated: list[Message] = []
+        aggregated_usage: UsageDetails | None = None
+        final_result: AgentResponse | None = None
+        while True:
+            await call_next()
+            iteration += 1
+
+            result = context.result
+            if not isinstance(result, AgentResponse):
+                raise TypeError(
+                    "AgentLoopMiddleware expected an AgentResponse from a non-streaming run, "
+                    f"got {type(result).__name__}."
+                )
+
+            final_result = result
+            aggregated.extend(result.messages)
+            if result.usage_details is not None:
+                aggregated_usage = add_usage_details(aggregated_usage, result.usage_details)
+
+            messages_used = context.messages
+            loop_kwargs = self._build_loop_kwargs(
+                context=context,
+                iteration=iteration,
+                last_result=result,
+                messages_used=messages_used,
+                original_messages=original_messages,
+                progress=progress,
+            )
+
+            work_iterations += 1
+            # Decide whether to stop and capture any feedback from should_continue first, so the
+            # feedback is available to both the progress and next-message callables this iteration.
+            stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations)
+            loop_kwargs = self._build_loop_kwargs(
+                context=context,
+                iteration=iteration,
+                last_result=result,
+                messages_used=messages_used,
+                original_messages=original_messages,
+                progress=progress,
+                feedback=feedback,
+            )
+            # Capture this iteration's progress entry, then refresh loop_kwargs so the next-message
+            # resolution sees the latest entry.
+            if await self._record_progress(result, loop_kwargs, progress):
+                loop_kwargs = self._build_loop_kwargs(
+                    context=context,
+                    iteration=iteration,
+                    last_result=result,
+                    messages_used=messages_used,
+                    original_messages=original_messages,
+                    progress=progress,
+                    feedback=feedback,
+                )
+            if stop:
+                break
+            if snapshot is not None and context.session is not None:
+                # Reset the session to the pre-loop baseline so the next run starts fresh; only the
+                # progress log (injected by _resolve_next_message) carries continuity forward.
+                self._restore_session(context.session, snapshot)
+            next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages)
+            context.messages = next_messages
+            aggregated.extend(next_messages)
+
+        if not self.return_final_only:
+            context.result = self._aggregate_response(final_result, aggregated, aggregated_usage)
+
+    def _process_streaming(
+        self,
+        context: AgentContext,
+        call_next: Callable[[], Awaitable[None]],
+        original_messages: list[Message],
+        snapshot: dict[str, Any] | None,
+    ) -> None:
+        # Holds the last iteration's final response so the outer stream's finalizer can return it
+        # rather than an aggregate of every iteration.
+        holder: dict[str, AgentResponse | None] = {"final": None}
+
+        async def _generator() -> Any:
+            iteration = 0
+            work_iterations = 0
+            progress: list[str] = []
+            while True:
+                try:
+                    await call_next()
+                    inner = context.result
+                    if not isinstance(inner, ResponseStream):
+                        raise TypeError(
+                            "AgentLoopMiddleware expected a ResponseStream from a streaming run, "
+                            f"got {type(inner).__name__}."
+                        )
+
+                    async for update in inner:
+                        yield update
+
+                    holder["final"] = await inner.get_final_response()
+                except MiddlewareTermination:
+                    # The pipeline's MiddlewareTermination suppression is no longer active once
+                    # process() has returned (the stream is consumed lazily), so a termination
+                    # raised by a downstream middleware or during stream consumption surfaces here.
+                    # Stop cleanly and keep whatever final response we have from a prior iteration.
+                    return
+
+                iteration += 1
+
+                messages_used = context.messages
+                final = holder["final"]
+                loop_kwargs = self._build_loop_kwargs(
+                    context=context,
+                    iteration=iteration,
+                    last_result=final,
+                    messages_used=messages_used,
+                    original_messages=original_messages,
+                    progress=progress,
+                )
+
+                work_iterations += 1
+                # Decide whether to stop and capture any feedback from should_continue first, so the
+                # feedback is available to both the progress and next-message callables this iteration.
+                stop, feedback = await self._evaluate_stop(loop_kwargs, work_iterations)
+                loop_kwargs = self._build_loop_kwargs(
+                    context=context,
+                    iteration=iteration,
+                    last_result=final,
+                    messages_used=messages_used,
+                    original_messages=original_messages,
+                    progress=progress,
+                    feedback=feedback,
+                )
+                if await self._record_progress(final, loop_kwargs, progress):
+                    loop_kwargs = self._build_loop_kwargs(
+                        context=context,
+                        iteration=iteration,
+                        last_result=final,
+                        messages_used=messages_used,
+                        original_messages=original_messages,
+                        progress=progress,
+                        feedback=feedback,
+                    )
+                if stop:
+                    return
+                if snapshot is not None and context.session is not None:
+                    # Reset the session to the pre-loop baseline before the next run. The final
+                    # response was already awaited above, so the service-side conversation id has
+                    # been propagated and is safe to discard here.
+                    self._restore_session(context.session, snapshot)
+                next_messages = await self._resolve_next_message(loop_kwargs, messages_used, original_messages)
+                context.messages = next_messages
+                # Surface the injected "nudge" messages in the stream so consumers see the user
+                # turns that drive each subsequent iteration (the equivalent of the aggregated
+                # transcript that non-streaming runs return).
+                for message in next_messages:
+                    yield self._message_to_update(message)
+
+        def _finalize(updates: Sequence[AgentResponseUpdate]) -> AgentResponse:
+            if holder["final"] is not None:
+                return holder["final"]
+            return AgentResponse.from_updates(updates)
+
+        context.result = ResponseStream(_generator(), finalizer=_finalize)
+
+    def _build_loop_kwargs(
+        self,
+        *,
+        context: AgentContext,
+        iteration: int,
+        last_result: AgentResponse | None,
+        messages_used: list[Message],
+        original_messages: list[Message],
+        progress: list[str],
+        feedback: str | None = None,
+    ) -> dict[str, Any]:
+        return {
+            "iteration": iteration,
+            "last_result": last_result,
+            "messages": messages_used,
+            "original_messages": original_messages,
+            "session": context.session,
+            "agent": context.agent,
+            # A copy so user callbacks cannot mutate the loop's internal progress log.
+            "progress": list(progress),
+            # Feedback returned by ``should_continue`` for this iteration (``None`` if it returned a
+            # plain bool, or the stop was decided by ``max_iterations``).
+            "feedback": feedback,
+        }
+
+    async def _record_progress(
+        self,
+        last_result: AgentResponse | None,
+        loop_kwargs: dict[str, Any],
+        progress: list[str],
+    ) -> bool:
+        """Capture this iteration's feedback into ``progress``. Returns ``True`` if an entry was added."""
+        if self.record_feedback is not None:
+            entry = await _maybe_await(self.record_feedback(**loop_kwargs))
+        else:
+            entry = last_result.text.strip() if last_result is not None else None
+        if entry:
+            progress.append(entry)
+            return True
+        return False
+
+    async def _evaluate_stop(self, loop_kwargs: dict[str, Any], work_iterations: int) -> tuple[bool, str | None]:
+        """Decide whether the loop should stop, returning ``(stop, feedback)``.
+
+        ``max_iterations`` is a safety cap that short-circuits before ``should_continue`` is
+        evaluated (so an expensive predicate/judge is not called once the cap has fired). Any
+        feedback returned by ``should_continue`` is propagated so the progress and next-message
+        callables can reference it.
+        """
+        if self.max_iterations is not None and work_iterations >= self.max_iterations:
+            return True, None
+        keep_going, feedback = await self._should_continue(loop_kwargs)
+        return (not keep_going), feedback
+
+    async def _should_continue(self, loop_kwargs: dict[str, Any]) -> tuple[bool, str | None]:
+        """Evaluate the predicate, normalizing its result to ``(continue, feedback)``."""
+        result = await _maybe_await(self.should_continue(**loop_kwargs))
+        return (bool(result[0]), result[1]) if isinstance(result, tuple) else (bool(result), None)  # type: ignore
+
+    @staticmethod
+    def _message_to_update(message: Message) -> AgentResponseUpdate:
+        """Wrap an injected loop message as a streaming update so consumers see it inline."""
+        return AgentResponseUpdate(
+            contents=message.contents,
+            role=message.role,
+            author_name=message.author_name,
+            message_id=message.message_id,
+        )
+
+    @staticmethod
+    def _aggregate_response(
+        final: AgentResponse,
+        messages: list[Message],
+        usage: UsageDetails | None,
+    ) -> AgentResponse:
+        """Build a combined response carrying every iteration's messages and summed usage.
+
+        Metadata (``response_id``, structured ``value``, etc.) is taken from the final iteration; the
+        structured value is passed through pre-parsed so it is not re-derived from the aggregated text.
+        """
+        return AgentResponse(
+            messages=messages,
+            response_id=final.response_id,
+            agent_id=final.agent_id,
+            created_at=final.created_at,
+            finish_reason=final.finish_reason,  # pyright: ignore[reportArgumentType]
+            usage_details=usage,
+            value=final.value,
+            additional_properties=dict(final.additional_properties) if final.additional_properties else None,
+            raw_representation=final.raw_representation,
+        )
+
+    @staticmethod
+    def _render_progress(entries: list[str]) -> Message:
+        """Format progress-log entries into a single ``user`` message."""
+        body = "\n".join(f"- {entry}" for entry in entries)
+        return Message(role="user", contents=[f"Progress so far:\n{body}"])
+
+    async def _resolve_next_message(
+        self,
+        loop_kwargs: dict[str, Any],
+        messages_used: list[Message],
+        original_messages: list[Message],
+    ) -> list[Message]:
+        # Compute the base next input. A ``next_message`` callable returning None requests a verbatim
+        # reuse of the previous messages (no progress injection); in fresh-context mode that escape
+        # hatch does not apply, so fall back to the default nudge instead.
+        if self.next_message is None:
+            next_msgs = normalize_messages(DEFAULT_NEXT_MESSAGE)
+        else:
+            next_input = await _maybe_await(self.next_message(**loop_kwargs))
+            if next_input is None:
+                if not self.fresh_context:
+                    return list(messages_used)
+                next_msgs = normalize_messages(DEFAULT_NEXT_MESSAGE)
+            else:
+                next_msgs = normalize_messages(next_input)
+
+        progress: list[str] = loop_kwargs.get("progress") or []
+        session = loop_kwargs.get("session")
+        progress_msg: Message | None = None
+        if self.inject_progress and progress:
+            # With a session the earlier entries are already retained in the conversation, so only
+            # the latest entry is injected to avoid duplication. Otherwise inject the full log.
+            entries = progress if (session is None or self.fresh_context) else progress[-1:]
+            progress_msg = self._render_progress(entries)
+
+        if self.fresh_context:
+            result = list(original_messages)
+            if progress_msg is not None:
+                result.append(progress_msg)
+            result.extend(next_msgs)
+            return result
+
+        if progress_msg is not None:
+            return [progress_msg, *next_msgs]
+        return list(next_msgs)
+
+
+def todos_remaining(provider: Any) -> ShouldContinueCallable:
+    """Build a ``should_continue`` predicate that loops while a ``TodoProvider`` has open items.
+
+    Args:
+        provider: A :class:`~agent_framework.TodoProvider` attached to the same session as the loop.
+
+    Returns:
+        A predicate suitable for :class:`AgentLoopMiddleware`'s ``should_continue`` argument.
+    """
+
+    async def _should_continue(*, session: Any = None, **kwargs: Any) -> bool:
+        if session is None:
+            return False
+        items = await provider.store.load_items(session, source_id=provider.source_id)
+        return any(not item.is_complete for item in items)
+
+    return _should_continue
+
+
+def background_tasks_running(provider: Any) -> ShouldContinueCallable:
+    """Build a ``should_continue`` predicate that loops while a ``BackgroundAgentsProvider`` is busy.
+
+    The predicate inspects the provider's persisted task state and continues while any task is still
+    marked as running. Pair it with ``max_iterations`` so the loop is guaranteed to stop even if a
+    task's persisted status is never refreshed.
+
+    Args:
+        provider: A :class:`~agent_framework.BackgroundAgentsProvider` attached to the same session
+            as the loop.
+
+    Returns:
+        A predicate suitable for :class:`AgentLoopMiddleware`'s ``should_continue`` argument.
+    """
+    from ._background_agents import BackgroundTaskInfo, BackgroundTaskStatus
+
+    def _should_continue(*, session: Any = None, **kwargs: Any) -> bool:
+        if session is None:
+            return False
+        state = session.state.get(provider.source_id)
+        if not state:
+            return False
+        return any(
+            BackgroundTaskInfo.from_dict(task).status == BackgroundTaskStatus.RUNNING for task in state.get("tasks", [])
+        )
+
+    return _should_continue

python/packages/core/agent_framework/_types.py

@@ -400,12 +400,18 @@ class UsageDetails(TypedDict, total=False, extra_items=int):  # type: ignore[cal
         input_token_count: The number of input tokens used.
         output_token_count: The number of output tokens generated.
         total_token_count: The total number of tokens (input + output).
+        cache_creation_input_token_count: The number of input tokens written to a provider-managed cache.
+        cache_read_input_token_count: The number of input tokens served from a provider-managed cache.
+        reasoning_output_token_count: The number of output tokens used for reasoning.
 
     """
 
     input_token_count: int | None
     output_token_count: int | None
     total_token_count: int | None
+    cache_creation_input_token_count: int | None
+    cache_read_input_token_count: int | None
+    reasoning_output_token_count: int | None
 
 
 def add_usage_details(usage1: UsageDetails | None, usage2: UsageDetails | None) -> UsageDetails:

python/packages/core/agent_framework/ag_ui/__init__.py

@@ -11,6 +11,10 @@ Supported classes and functions:
 - AGUIChatClient
 - AGUIEventConverter
 - AGUIHttpService
+- AGUIThreadSnapshot
+- AGUIThreadSnapshotStore
+- InMemoryAGUIThreadSnapshotStore
+- SnapshotScopeResolver
 - add_agent_framework_fastapi_endpoint
 - state_update
 - __version__
@@ -28,6 +32,10 @@ _IMPORTS = [
     "AGUIChatClient",
     "AGUIEventConverter",
     "AGUIHttpService",
+    "AGUIThreadSnapshot",
+    "AGUIThreadSnapshotStore",
+    "InMemoryAGUIThreadSnapshotStore",
+    "SnapshotScopeResolver",
     "state_update",
     "__version__",
 ]

python/packages/core/agent_framework/ag_ui/__init__.pyi

@@ -6,6 +6,10 @@ from agent_framework_ag_ui import (
     AGUIChatClient,
     AGUIEventConverter,
     AGUIHttpService,
+    AGUIThreadSnapshot,
+    AGUIThreadSnapshotStore,
+    InMemoryAGUIThreadSnapshotStore,
+    SnapshotScopeResolver,
     __version__,
     add_agent_framework_fastapi_endpoint,
     state_update,
@@ -15,8 +19,12 @@ __all__ = [
     "AGUIChatClient",
     "AGUIEventConverter",
     "AGUIHttpService",
+    "AGUIThreadSnapshot",
+    "AGUIThreadSnapshotStore",
     "AgentFrameworkAgent",
     "AgentFrameworkWorkflow",
+    "InMemoryAGUIThreadSnapshotStore",
+    "SnapshotScopeResolver",
     "__version__",
     "add_agent_framework_fastapi_endpoint",
     "state_update",

python/packages/core/agent_framework/observability.py

@@ -201,6 +201,9 @@ class OtelAttr(str, Enum):
     # Usage attributes
     INPUT_TOKENS = "gen_ai.usage.input_tokens"
     OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+    CACHE_CREATION_INPUT_TOKENS = "gen_ai.usage.cache_creation.input_tokens"
+    CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read.input_tokens"
+    REASONING_OUTPUT_TOKENS = "gen_ai.usage.reasoning.output_tokens"
     # Tool attributes
     TOOL_CALL_ID = "gen_ai.tool.call.id"
     TOOL_DESCRIPTION = "gen_ai.tool.description"
@@ -327,6 +330,20 @@ FINISH_REASON_MAP = {
     "tool_calls": "tool_call",
     "length": "length",
 }
+USAGE_DETAIL_TO_OTEL_ATTR: Final[tuple[tuple[str, OtelAttr], ...]] = (
+    ("input_token_count", OtelAttr.INPUT_TOKENS),
+    ("output_token_count", OtelAttr.OUTPUT_TOKENS),
+    ("cache_creation_input_token_count", OtelAttr.CACHE_CREATION_INPUT_TOKENS),
+    ("cache_read_input_token_count", OtelAttr.CACHE_READ_INPUT_TOKENS),
+    ("reasoning_output_token_count", OtelAttr.REASONING_OUTPUT_TOKENS),
+    ("anthropic.cache_creation_input_tokens", OtelAttr.CACHE_CREATION_INPUT_TOKENS),
+    ("anthropic.cache_read_input_tokens", OtelAttr.CACHE_READ_INPUT_TOKENS),
+    ("openai.cached_input_tokens", OtelAttr.CACHE_READ_INPUT_TOKENS),
+    ("prompt/cached_tokens", OtelAttr.CACHE_READ_INPUT_TOKENS),
+    ("openai.reasoning_tokens", OtelAttr.REASONING_OUTPUT_TOKENS),
+    ("completion/reasoning_tokens", OtelAttr.REASONING_OUTPUT_TOKENS),
+    ("reasoning_tokens", OtelAttr.REASONING_OUTPUT_TOKENS),
+)
 
 
 # region Telemetry utils
@@ -2350,12 +2367,16 @@ def _apply_accumulated_usage(attributes: dict[str, Any], captured_fields: set[st
     accumulated = INNER_ACCUMULATED_USAGE.get()
     if not accumulated:
         return
-    input_tokens = accumulated.get("input_token_count")
-    if input_tokens:
-        attributes[OtelAttr.INPUT_TOKENS] = input_tokens
-    output_tokens = accumulated.get("output_token_count")
-    if output_tokens:
-        attributes[OtelAttr.OUTPUT_TOKENS] = output_tokens
+    _apply_usage_attributes(attributes, accumulated)
+
+
+def _apply_usage_attributes(attributes: dict[str, Any], usage: Mapping[str, Any]) -> None:
+    """Apply known usage details as standard OTel GenAI attributes."""
+    for usage_key, otel_attr in USAGE_DETAIL_TO_OTEL_ATTR:
+        value = usage.get(usage_key)
+        if value is None or isinstance(value, bool) or not isinstance(value, int):
+            continue
+        attributes.setdefault(otel_attr, value)
 
 
 def _get_response_attributes(
@@ -2378,12 +2399,7 @@ def _get_response_attributes(
     if model := getattr(response, "model", None):
         attributes[OtelAttr.RESPONSE_MODEL] = model
     if capture_usage and (usage := response.usage_details):
-        input_tokens = usage.get("input_token_count")
-        if input_tokens:
-            attributes[OtelAttr.INPUT_TOKENS] = input_tokens
-        output_tokens = usage.get("output_token_count")
-        if output_tokens:
-            attributes[OtelAttr.OUTPUT_TOKENS] = output_tokens
+        _apply_usage_attributes(attributes, usage)
     return attributes
 
 
@@ -2407,9 +2423,9 @@ def _capture_response(
     """Set the response for a given span."""
     span.set_attributes(attributes)
     attrs: dict[str, Any] = {k: v for k, v in attributes.items() if k in GEN_AI_METRIC_ATTRIBUTES}
-    if token_usage_histogram and (input_tokens := attributes.get(OtelAttr.INPUT_TOKENS)):
+    if token_usage_histogram and (input_tokens := attributes.get(OtelAttr.INPUT_TOKENS)) is not None:
         token_usage_histogram.record(input_tokens, attributes={**attrs, OtelAttr.T_TYPE: OtelAttr.T_TYPE_INPUT})
-    if token_usage_histogram and (output_tokens := attributes.get(OtelAttr.OUTPUT_TOKENS)):
+    if token_usage_histogram and (output_tokens := attributes.get(OtelAttr.OUTPUT_TOKENS)) is not None:
         token_usage_histogram.record(output_tokens, {**attrs, OtelAttr.T_TYPE: OtelAttr.T_TYPE_OUTPUT})
     if operation_duration_histogram and duration is not None:
         if OtelAttr.ERROR_TYPE in attributes:

python/packages/core/tests/core/test_harness_file_access.py

@@ -158,6 +158,69 @@ async def test_in_memory_store_search_returns_matches_with_snippets() -> None:
     assert {result.file_name for result in results_all} == {"a.md", "notes.txt"}
 
 
+async def test_in_memory_store_search_is_recursive_with_root_relative_names() -> None:
+    """Recursive search should find files at any depth and return root-relative names."""
+    store = InMemoryAgentFileStore()
+    await store.write_file("top.md", "ERROR at top")
+    await store.write_file("reports/q1.md", "ERROR in q1")
+    await store.write_file("reports/2024/q2.md", "ERROR in q2")
+    await store.write_file("reports/2024/data.txt", "ERROR wrong extension")
+
+    # Non-recursive (default) only sees the direct child.
+    direct = await store.search_files("", "error")
+    assert {result.file_name for result in direct} == {"top.md"}
+
+    # Recursive sees every descendant, with store-root-relative file names.
+    recursive = await store.search_files("", "error", recursive=True)
+    assert {result.file_name for result in recursive} == {
+        "top.md",
+        "reports/q1.md",
+        "reports/2024/q2.md",
+        "reports/2024/data.txt",
+    }
+
+    # Subtree scoping via the glob (``*`` crosses ``/`` with fnmatch).
+    scoped = await store.search_files("", "error", "reports/*", recursive=True)
+    assert {result.file_name for result in scoped} == {
+        "reports/q1.md",
+        "reports/2024/q2.md",
+        "reports/2024/data.txt",
+    }
+
+    # Extension glob matches markdown at any depth but not other extensions.
+    markdown = await store.search_files("", "error", "*.md", recursive=True)
+    assert {result.file_name for result in markdown} == {
+        "top.md",
+        "reports/q1.md",
+        "reports/2024/q2.md",
+    }
+
+
+async def test_in_memory_store_list_directories() -> None:
+    """``list_directories`` should return direct child subdirectories only, preserving casing."""
+    store = InMemoryAgentFileStore()
+    await store.write_file("top.md", "x")
+    await store.write_file("Reports/q1.md", "x")
+    await store.write_file("Reports/2024/q2.md", "x")
+    await store.write_file("data/raw.csv", "x")
+
+    assert sorted(await store.list_directories()) == ["Reports", "data"]
+    assert sorted(await store.list_directories("Reports")) == ["2024"]
+    # A directory with no subdirectories returns an empty list.
+    assert await store.list_directories("data") == []
+    # A missing directory returns an empty list.
+    assert await store.list_directories("missing") == []
+
+
+async def test_in_memory_store_list_directories_rejects_traversal() -> None:
+    """``list_directories`` must reject traversal inputs the way ``list_files`` does."""
+    store = InMemoryAgentFileStore()
+    await store.write_file("reports/q1.md", "x")
+    for bad in ("../escape", "/abs/path", ".."):
+        with pytest.raises(ValueError):
+            await store.list_directories(bad)
+
+
 async def test_in_memory_store_search_rejects_invalid_and_oversize_regex() -> None:
     """``search_files`` should surface clean errors for bad regex input."""
     store = InMemoryAgentFileStore()
@@ -267,6 +330,78 @@ async def test_filesystem_store_search_matches_lines_and_filters_globs(tmp_path:
     assert {result.file_name for result in results_all} == {"a.md", "b.txt"}
 
 
+async def test_filesystem_store_search_is_recursive_with_root_relative_names(tmp_path: Path) -> None:
+    """Recursive filesystem search should walk the subtree and return root-relative names."""
+    store = FileSystemAgentFileStore(tmp_path)
+    await store.write_file("top.md", "ERROR at top")
+    await store.write_file("reports/q1.md", "ERROR in q1")
+    await store.write_file("reports/2024/q2.md", "ERROR in q2")
+
+    direct = await store.search_files("", "error")
+    assert {result.file_name for result in direct} == {"top.md"}
+
+    recursive = await store.search_files("", "error", recursive=True)
+    assert {result.file_name for result in recursive} == {
+        "top.md",
+        "reports/q1.md",
+        "reports/2024/q2.md",
+    }
+
+    scoped = await store.search_files("", "error", "reports/*", recursive=True)
+    assert {result.file_name for result in scoped} == {
+        "reports/q1.md",
+        "reports/2024/q2.md",
+    }
+
+
+async def test_filesystem_store_list_directories(tmp_path: Path) -> None:
+    """``list_directories`` should list direct child subdirectories only."""
+    store = FileSystemAgentFileStore(tmp_path)
+    await store.write_file("top.md", "x")
+    await store.write_file("reports/q1.md", "x")
+    await store.write_file("reports/2024/q2.md", "x")
+    await store.write_file("data/raw.csv", "x")
+
+    assert sorted(await store.list_directories()) == ["data", "reports"]
+    assert sorted(await store.list_directories("reports")) == ["2024"]
+    assert await store.list_directories("data") == []
+    assert await store.list_directories("missing") == []
+
+
+async def test_filesystem_store_list_directories_rejects_traversal(tmp_path: Path) -> None:
+    """``list_directories`` is security-critical and must reject paths that escape the root."""
+    store = FileSystemAgentFileStore(tmp_path)
+    await store.write_file("reports/q1.md", "x")
+    for bad in ("../escape", "/etc", "C:/Windows", ".."):
+        with pytest.raises(ValueError):
+            await store.list_directories(bad)
+
+
+async def test_filesystem_store_search_and_list_skip_symlinked_directories(tmp_path: Path) -> None:
+    """Recursive search must not descend into symlinked dirs and ``list_directories`` must exclude them."""
+    target = tmp_path / "outside"
+    target.mkdir()
+    (target / "secret.md").write_text("ERROR outside the root", encoding="utf-8")
+
+    root = tmp_path / "root"
+    root.mkdir()
+    (root / "inside.md").write_text("ERROR inside", encoding="utf-8")
+    link = root / "linked"
+    try:
+        link.symlink_to(target, target_is_directory=True)
+    except (OSError, NotImplementedError) as exc:
+        pytest.skip(f"Symbolic links are not supported in this environment: {exc!r}")
+
+    store = FileSystemAgentFileStore(root)
+
+    # ``list_directories`` excludes the symlinked directory.
+    assert await store.list_directories() == []
+
+    # Recursive search does not follow the symlink out of the root.
+    results = await store.search_files("", "error", recursive=True)
+    assert {result.file_name for result in results} == {"inside.md"}
+
+
 async def test_filesystem_store_search_skips_non_utf8_files(tmp_path: Path) -> None:
     """The filesystem store should silently skip non-UTF-8 files instead of aborting the search."""
     store = FileSystemAgentFileStore(tmp_path)
@@ -303,7 +438,7 @@ def test_filesystem_store_requires_non_empty_root() -> None:
 async def test_file_access_provider_registers_tools_and_instructions(
     chat_client_base: SupportsChatGetResponse,
 ) -> None:
-    """``FileAccessProvider.before_run`` should add the canonical instructions and five tools."""
+    """``FileAccessProvider.before_run`` should add the canonical instructions and six tools."""
     session = AgentSession(session_id="session-1")
     store = InMemoryAgentFileStore()
     provider = FileAccessProvider(store=store)
@@ -321,6 +456,7 @@ async def test_file_access_provider_registers_tools_and_instructions(
         "file_access_read_file",
         "file_access_delete_file",
         "file_access_list_files",
+        "file_access_list_subdirectories",
         "file_access_search_files",
     }
     assert {getattr(tool, "name", None) for tool in tools} >= expected_names
@@ -354,6 +490,7 @@ async def test_file_access_provider_delete_approval_defaults_to_always_require(
         "file_access_save_file",
         "file_access_read_file",
         "file_access_list_files",
+        "file_access_list_subdirectories",
         "file_access_search_files",
     ):
         assert _tool_by_name(tools, name).approval_mode == "never_require"
@@ -396,6 +533,7 @@ async def test_file_access_provider_tools_round_trip_files(
     read_file = _tool_by_name(tools, "file_access_read_file")
     delete_file = _tool_by_name(tools, "file_access_delete_file")
     list_files = _tool_by_name(tools, "file_access_list_files")
+    list_subdirectories = _tool_by_name(tools, "file_access_list_subdirectories")
     search_files = _tool_by_name(tools, "file_access_search_files")
 
     saved = await save_file.invoke(arguments={"file_name": "plan.md", "content": "step 1\nERROR step 2"})
@@ -426,6 +564,15 @@ async def test_file_access_provider_tools_round_trip_files(
     listed_blank = await list_files.invoke(arguments={"directory": "   "})
     assert sorted(json.loads(listed_blank[0].text)) == ["plan.md"]
 
+    # The subdirectory-discovery tool surfaces child directories (not files).
+    listed_dirs = await list_subdirectories.invoke()
+    assert json.loads(listed_dirs[0].text) == ["reports"]
+    listed_dirs_blank = await list_subdirectories.invoke(arguments={"directory": "   "})
+    assert json.loads(listed_dirs_blank[0].text) == ["reports"]
+    # A leaf directory with no child directories returns an empty list.
+    listed_dirs_nested = await list_subdirectories.invoke(arguments={"directory": "reports"})
+    assert json.loads(listed_dirs_nested[0].text) == []
+
     missing = await read_file.invoke(arguments={"file_name": "missing.md"})
     assert "not found" in missing[0].text
 
@@ -434,14 +581,12 @@ async def test_file_access_provider_tools_round_trip_files(
     assert parsed[0]["file_name"] == "plan.md"
     assert parsed[0]["matching_lines"][0]["line"] == "ERROR replaced"
 
-    # The search tool should likewise accept an optional directory argument so
-    # agents can scope a search to a subfolder.
+    # The search tool is recursive from the store root; scope to a subtree using
+    # the glob (``*`` crosses ``/`` with fnmatch). Results use root-relative names.
     await save_file.invoke(arguments={"file_name": "reports/issues.md", "content": "ERROR nested"})
-    scoped = await search_files.invoke(
-        arguments={"regex_pattern": "error", "file_pattern": "*.md", "directory": "reports"}
-    )
+    scoped = await search_files.invoke(arguments={"regex_pattern": "error", "file_pattern": "reports/*"})
     scoped_parsed = json.loads(scoped[0].text)
-    assert [entry["file_name"] for entry in scoped_parsed] == ["issues.md"]
+    assert [entry["file_name"] for entry in scoped_parsed] == ["reports/issues.md"]
 
     deleted = await delete_file.invoke(arguments={"file_name": "plan.md"})
     assert "deleted" in deleted[0].text

python/packages/core/tests/core/test_observability.py

@@ -2154,6 +2154,58 @@ def test_get_response_attributes_with_usage():
     assert result[OtelAttr.OUTPUT_TOKENS] == 50
 
 
+def test_get_response_attributes_with_additional_usage():
+    """Test _get_response_attributes maps additional usage details to OTel attributes."""
+    from unittest.mock import Mock
+
+    from agent_framework.observability import OtelAttr, _get_response_attributes
+
+    response = Mock()
+    response.response_id = None
+    response.finish_reason = None
+    response.raw_representation = None
+    response.usage_details = {
+        "input_token_count": 0,
+        "output_token_count": 50,
+        "cache_creation_input_token_count": 10,
+        "cache_read_input_token_count": 0,
+        "reasoning_output_token_count": 30,
+    }
+
+    attrs = {}
+    result = _get_response_attributes(attrs, response)
+
+    assert result[OtelAttr.INPUT_TOKENS] == 0
+    assert result[OtelAttr.OUTPUT_TOKENS] == 50
+    assert result[OtelAttr.CACHE_CREATION_INPUT_TOKENS] == 10
+    assert result[OtelAttr.CACHE_READ_INPUT_TOKENS] == 0
+    assert result[OtelAttr.REASONING_OUTPUT_TOKENS] == 30
+
+
+def test_get_response_attributes_maps_legacy_usage_keys():
+    """Test _get_response_attributes maps legacy provider usage keys to standard OTel attributes."""
+    from unittest.mock import Mock
+
+    from agent_framework.observability import OtelAttr, _get_response_attributes
+
+    response = Mock()
+    response.response_id = None
+    response.finish_reason = None
+    response.raw_representation = None
+    response.usage_details = {
+        "anthropic.cache_creation_input_tokens": 12,
+        "openai.cached_input_tokens": 0,
+        "completion/reasoning_tokens": 34,
+    }
+
+    attrs = {}
+    result = _get_response_attributes(attrs, response)
+
+    assert result[OtelAttr.CACHE_CREATION_INPUT_TOKENS] == 12
+    assert result[OtelAttr.CACHE_READ_INPUT_TOKENS] == 0
+    assert result[OtelAttr.REASONING_OUTPUT_TOKENS] == 34
+
+
 def test_get_response_attributes_capture_usage_false():
     """Test _get_response_attributes skips usage when capture_usage is False."""
     from unittest.mock import Mock
@@ -2164,13 +2216,22 @@ def test_get_response_attributes_capture_usage_false():
     response.response_id = None
     response.finish_reason = None
     response.raw_representation = None
-    response.usage_details = {"input_token_count": 100, "output_token_count": 50}
+    response.usage_details = {
+        "input_token_count": 100,
+        "output_token_count": 50,
+        "cache_creation_input_token_count": 10,
+        "cache_read_input_token_count": 20,
+        "reasoning_output_token_count": 30,
+    }
 
     attrs = {}
     result = _get_response_attributes(attrs, response, capture_usage=False)
 
     assert OtelAttr.INPUT_TOKENS not in result
     assert OtelAttr.OUTPUT_TOKENS not in result
+    assert OtelAttr.CACHE_CREATION_INPUT_TOKENS not in result
+    assert OtelAttr.CACHE_READ_INPUT_TOKENS not in result
+    assert OtelAttr.REASONING_OUTPUT_TOKENS not in result
 
 
 def test_get_response_attributes_capture_response_id_false():
@@ -2933,6 +2994,23 @@ def test_capture_response(span_exporter: InMemorySpanExporter):
     assert spans[0].attributes.get(OtelAttr.OUTPUT_TOKENS) == 50
 
 
+def test_capture_response_records_zero_token_usage():
+    """Test _capture_response records zero-valued token usage."""
+    from agent_framework.observability import OtelAttr, _capture_response
+
+    span = Mock()
+    token_histogram = Mock()
+    attrs = {
+        OtelAttr.INPUT_TOKENS: 0,
+        OtelAttr.OUTPUT_TOKENS: 0,
+    }
+
+    _capture_response(span=span, attributes=attrs, token_usage_histogram=token_histogram)
+
+    span.set_attributes.assert_called_once_with(attrs)
+    assert token_histogram.record.call_count == 2
+
+
 async def test_layer_ordering_span_sequence_with_function_calling(span_exporter: InMemorySpanExporter):
     """Test that with correct layer ordering, spans appear in the expected sequence.
 
@@ -3937,11 +4015,21 @@ async def test_agent_invoke_span_aggregates_usage_across_tool_calls(span_exporte
                     Content.from_function_call(call_id="call_1", name="get_weather", arguments='{"city": "Seattle"}')
                 ],
             ),
-            usage_details=UsageDetails(input_token_count=2239, output_token_count=192),
+            usage_details=UsageDetails(
+                input_token_count=2239,
+                output_token_count=192,
+                cache_read_input_token_count=100,
+                reasoning_output_token_count=25,
+            ),
         ),
         ChatResponse(
             messages=Message(role="assistant", contents=["The weather in Seattle is sunny."]),
-            usage_details=UsageDetails(input_token_count=2569, output_token_count=99),
+            usage_details=UsageDetails(
+                input_token_count=2569,
+                output_token_count=99,
+                cache_read_input_token_count=200,
+                reasoning_output_token_count=0,
+            ),
         ),
     ]
 
@@ -3965,12 +4053,18 @@ async def test_agent_invoke_span_aggregates_usage_across_tool_calls(span_exporte
     # Individual chat spans retain their own usage
     assert chat_spans[0].attributes.get(OtelAttr.INPUT_TOKENS) == 2239
     assert chat_spans[0].attributes.get(OtelAttr.OUTPUT_TOKENS) == 192
+    assert chat_spans[0].attributes.get(OtelAttr.CACHE_READ_INPUT_TOKENS) == 100
+    assert chat_spans[0].attributes.get(OtelAttr.REASONING_OUTPUT_TOKENS) == 25
     assert chat_spans[1].attributes.get(OtelAttr.INPUT_TOKENS) == 2569
     assert chat_spans[1].attributes.get(OtelAttr.OUTPUT_TOKENS) == 99
+    assert chat_spans[1].attributes.get(OtelAttr.CACHE_READ_INPUT_TOKENS) == 200
+    assert chat_spans[1].attributes.get(OtelAttr.REASONING_OUTPUT_TOKENS) == 0
 
     # The invoke_agent span must report the aggregate across all LLM round-trips
     assert agent_span.attributes.get(OtelAttr.INPUT_TOKENS) == 2239 + 2569
     assert agent_span.attributes.get(OtelAttr.OUTPUT_TOKENS) == 192 + 99
+    assert agent_span.attributes.get(OtelAttr.CACHE_READ_INPUT_TOKENS) == 100 + 200
+    assert agent_span.attributes.get(OtelAttr.REASONING_OUTPUT_TOKENS) == 25
 
 
 @pytest.mark.parametrize("enable_sensitive_data", [False], indirect=True)