Compare commits

..
73 changed files with 257 additions and 6948 deletions
+1 -4
View File
@@ -20,10 +20,7 @@ ignorePatterns:
- pattern: "https://your-resource.openai.azure.com/"
- pattern: "http://host.docker.internal"
- pattern: "https://openai.github.io/openai-agents-js/openai/agents/classes/"
# dotnet.microsoft.com bot-blocks CI link checkers with intermittent 403s on any
# path (including localized variants like /en-us/download/...), so ignore the
# whole domain rather than just /download.
- pattern: "https:\/\/dotnet.microsoft.com"
- pattern: "https:\/\/dotnet.microsoft.com\/download"
- pattern: "https://github.com/Rel1cx/eslint-react"
# excludedDirs:
# Folders which include links to localhost, since it's not ignored with regular expressions
+5 -25
View File
@@ -1,43 +1,23 @@
### Motivation & Context
### Motivation and Context
<!-- Thank you for your contribution to the Agent Framework repo!
Please help reviewers and future users, providing the following information:
1. Why is this change required?
2. What problem does it solve?
3. What scenario does it contribute to?
4. If it fixes an open issue, please link to the issue below.
4. If it fixes an open issue, please link to the issue here.
-->
### Description & Review Guide
### Description
<!-- Describe your changes, the overall approach, the underlying design.
Highlight what you want the reviewers to focus on.
These notes will help understanding how your code works. Thanks! -->
- **What are the major changes?**
- **What is the impact of these changes?**
- **What do you want reviewers to focus on?**
<!-- NOTE FOR AUTOMATED/AI REVIEWERS: the "What do you want reviewers to focus on?"
item above is intended for human reviewers only. Automated/AI reviewers should
ignore it and review the entire change rather than narrowing scope to it. -->
### Related Issue
<!-- Which issue does this PR fix? Link it using a GitHub closing keyword so it is
closed automatically when this PR is merged, e.g. "Fixes #123" or "Closes #123".
PRs that are not linked to an issue may be closed, no matter how valid the change is.
Also check whether an open PR already exists for this issue; if so,
explain how this PR is different. -->
Fixes #
### Contribution Checklist
<!-- Before submitting this PR, please make sure: -->
- [ ] The code builds clean without any errors or warnings
- [ ] All unit tests pass, and I have added new tests where possible
- [ ] The PR follows the [Contribution Guidelines](https://github.com/microsoft/agent-framework/blob/main/CONTRIBUTING.md)
- [ ] This PR is linked to an issue and there is no other open PR for this issue (see Related Issue above).
- [x] **This is not a breaking change.** If it _is_ a breaking change, add the `breaking change` label (or add "[BREAKING]" to the title prefix, before or after any language prefix) — a workflow keeps the label and title prefix in sync automatically.
- [ ] All unit tests pass, and I have added new tests where possible
- [ ] **Is this a breaking change?** If yes, add "[BREAKING]" prefix to the title of the PR.
-253
View File
@@ -1,253 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
const BREAKING_CHANGE_LABEL = 'breaking change';
const BREAKING_PREFIX = '[BREAKING]';
const DEFAULT_PREFIX_LABELS = Object.freeze({
python: 'Python',
'.NET': '.NET',
});
const DEFAULT_BRACKET_PREFIX_LABELS = Object.freeze({
[BREAKING_CHANGE_LABEL]: BREAKING_PREFIX,
});
function escapeRegExp(value) {
return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
}
function getMatchingValueByKey(valuesByKey, keyToFind) {
const matchingKey = Object.keys(valuesByKey).find((key) => key.toLowerCase() === keyToFind.toLowerCase());
return matchingKey === undefined ? null : valuesByKey[matchingKey];
}
function getPrefixPattern(prefixes) {
return prefixes.map(escapeRegExp).join('|');
}
function canonicalizePrefix(prefix, prefixes) {
return prefixes.find((knownPrefix) => knownPrefix.toLowerCase() === prefix.toLowerCase()) ?? prefix;
}
function normalizeLeadingBracketPrefix(title, bracketPrefixes) {
const bracketPattern = getPrefixPattern(bracketPrefixes);
if (!bracketPattern) {
return title;
}
const leadingBracketPrefix = new RegExp(`^(${bracketPattern})(?=\\s|$)`, 'i');
return title.replace(
leadingBracketPrefix,
(bracketPrefix) => canonicalizePrefix(bracketPrefix, bracketPrefixes),
);
}
function parseLeadingTitlePrefix(title, titlePrefixes) {
const titlePrefixPattern = getPrefixPattern(titlePrefixes);
if (!titlePrefixPattern) {
return null;
}
const match = title.match(new RegExp(`^(${titlePrefixPattern}):\\s*`, 'i'));
if (!match) {
return null;
}
return {
prefix: canonicalizePrefix(match[1], titlePrefixes),
rest: title.slice(match[0].length).trimStart(),
};
}
function removeBracketPrefixToken(title, bracketPrefix) {
const bracketPrefixPattern = escapeRegExp(bracketPrefix);
return title
.replace(new RegExp(`(^|\\s+)${bracketPrefixPattern}(?=\\s|$)`, 'ig'), '$1')
.replace(/\s{2,}/g, ' ')
.trim();
}
function addTitlePrefix(title, prefix, bracketPrefixes = Object.values(DEFAULT_BRACKET_PREFIX_LABELS)) {
const bracketPattern = getPrefixPattern(bracketPrefixes);
const prefixPattern = escapeRegExp(prefix);
if (bracketPattern) {
const bracketThenTitlePrefix = new RegExp(`^(${bracketPattern})(\\s+)(${prefixPattern})(?=:)`, 'i');
if (bracketThenTitlePrefix.test(title)) {
return title.replace(
bracketThenTitlePrefix,
(match, bracketPrefix, spacing) => `${canonicalizePrefix(bracketPrefix, bracketPrefixes)}${spacing}${prefix}`,
);
}
title = normalizeLeadingBracketPrefix(title, bracketPrefixes);
}
if (!title.startsWith(`${prefix}: `)) {
const existingTitlePrefix = new RegExp(`^${prefixPattern}:\\s*`, 'i');
if (existingTitlePrefix.test(title)) {
return title.replace(existingTitlePrefix, `${prefix}: `);
}
return `${prefix}: ${title}`;
}
return title;
}
function hasBracketPrefix(title, bracketPrefix, titlePrefixes = Object.values(DEFAULT_PREFIX_LABELS)) {
const bracketPrefixPattern = escapeRegExp(bracketPrefix);
const leadingBracketPrefix = new RegExp(`^${bracketPrefixPattern}(?=\\s|$)`, 'i');
if (leadingBracketPrefix.test(title)) {
return true;
}
const leadingTitlePrefix = parseLeadingTitlePrefix(title, titlePrefixes);
if (!leadingTitlePrefix) {
return false;
}
return leadingBracketPrefix.test(leadingTitlePrefix.rest);
}
function addBracketPrefix(title, bracketPrefix, titlePrefixes = Object.values(DEFAULT_PREFIX_LABELS)) {
const bracketPrefixPattern = escapeRegExp(bracketPrefix);
const leadingBracketPrefix = new RegExp(`^${bracketPrefixPattern}(?=\\s|$)`, 'i');
if (leadingBracketPrefix.test(title)) {
return title.replace(leadingBracketPrefix, bracketPrefix);
}
const leadingTitlePrefix = parseLeadingTitlePrefix(title, titlePrefixes);
if (leadingTitlePrefix) {
if (leadingBracketPrefix.test(leadingTitlePrefix.rest)) {
const normalizedRest = leadingTitlePrefix.rest.replace(leadingBracketPrefix, bracketPrefix);
return `${leadingTitlePrefix.prefix}: ${normalizedRest}`;
}
const titleWithoutBracketPrefix = removeBracketPrefixToken(leadingTitlePrefix.rest, bracketPrefix);
return `${leadingTitlePrefix.prefix}: ${bracketPrefix}`
+ (titleWithoutBracketPrefix ? ` ${titleWithoutBracketPrefix}` : '');
}
const titleWithoutBracketPrefix = removeBracketPrefixToken(title, bracketPrefix);
return `${bracketPrefix}${titleWithoutBracketPrefix ? ` ${titleWithoutBracketPrefix}` : ''}`;
}
function hasLabel(labels, labelName) {
return labels.some((label) => label.toLowerCase() === labelName.toLowerCase());
}
function getCurrentTitle(context) {
switch (context.eventName) {
case 'issues':
return context.payload.issue.title;
case 'pull_request_target':
return context.payload.pull_request.title;
default:
throw new Error(`Unrecognized eventName: ${context.eventName}`);
}
}
async function updateTitleForAddedLabel({
github,
context,
core,
prefixLabels = DEFAULT_PREFIX_LABELS,
bracketPrefixLabels = DEFAULT_BRACKET_PREFIX_LABELS,
}) {
const labelAdded = context.payload.label?.name;
if (!labelAdded) {
throw new Error('This script must be run from a labeled event.');
}
const currentTitle = getCurrentTitle(context);
let newTitle = null;
const titlePrefix = getMatchingValueByKey(prefixLabels, labelAdded);
if (titlePrefix !== null) {
newTitle = addTitlePrefix(currentTitle, titlePrefix, Object.values(bracketPrefixLabels));
}
const bracketPrefix = getMatchingValueByKey(bracketPrefixLabels, labelAdded);
if (bracketPrefix !== null) {
newTitle = addBracketPrefix(currentTitle, bracketPrefix, Object.values(prefixLabels));
}
if (newTitle === null) {
core.info(`No title prefix configured for label "${labelAdded}".`);
return { updated: false, newTitle: currentTitle };
}
if (newTitle === currentTitle) {
core.info(`Title already includes the prefix for label "${labelAdded}".`);
return { updated: false, newTitle };
}
switch (context.eventName) {
case 'issues':
await github.rest.issues.update({
issue_number: context.issue.number,
owner: context.repo.owner,
repo: context.repo.repo,
title: newTitle,
});
break;
case 'pull_request_target':
await github.rest.pulls.update({
pull_number: context.issue.number,
owner: context.repo.owner,
repo: context.repo.repo,
title: newTitle,
});
break;
default:
throw new Error(`Unrecognized eventName: ${context.eventName}`);
}
return { updated: true, newTitle };
}
async function syncBreakingChangeLabelFromTitle({
github,
context,
core,
labelName = BREAKING_CHANGE_LABEL,
bracketPrefix = BREAKING_PREFIX,
titlePrefixes = Object.values(DEFAULT_PREFIX_LABELS),
}) {
const pullRequest = context.payload.pull_request;
if (!pullRequest) {
throw new Error('This script must be run from a pull_request_target event.');
}
const title = pullRequest.title || '';
if (!hasBracketPrefix(title, bracketPrefix, titlePrefixes)) {
core.info(`Title does not include ${bracketPrefix} in the title prefix.`);
return { added: false };
}
const labels = pullRequest.labels?.map((label) => label.name).filter(Boolean) ?? [];
if (hasLabel(labels, labelName)) {
core.info(`PR already has the "${labelName}" label.`);
return { added: false };
}
await github.rest.issues.addLabels({
issue_number: context.issue.number,
owner: context.repo.owner,
repo: context.repo.repo,
labels: [labelName],
});
return { added: true };
}
module.exports = {
addBracketPrefix,
addTitlePrefix,
hasBracketPrefix,
syncBreakingChangeLabelFromTitle,
updateTitleForAddedLabel,
};
-116
View File
@@ -1,116 +0,0 @@
---
name: pull-requests
description: >
Guidance for creating pull requests and handling PR review comments in the
Agent Framework repository. Use this when writing a PR description (filling out
the PR template) or when responding to and resolving review comments on an
existing PR.
---
# Pull Request Workflow
This skill covers two tasks: (1) writing a high-quality PR description, and
(2) handling review comments on an existing PR.
## 1. Writing the PR description
Always follow the repository PR template at
[`.github/pull_request_template.md`](../../pull_request_template.md). Keep its
exact structure and headings. Fill every section:
### `### Motivation & Context`
Explain *why* the change is needed: the problem it solves and the scenario it
contributes to. Describe the net change relative to `main` — this is implied, so
do **not** spell out "vs main" explicitly.
### `### Description & Review Guide`
Describe the changes, the overall approach, and the design. Answer the three
prompts:
- **What are the major changes?**
- **What is the impact of these changes?**
- **What do you want reviewers to focus on?** — This item is for **human
reviewers only**. Automated/AI reviewers must ignore it and review the entire
change rather than narrowing scope to it.
### `### Related Issue`
Link the issue the PR fixes using a GitHub closing keyword (`Fixes #123` /
`Closes #123`) so it closes automatically on merge. A PR with no linked issue may
be closed regardless of how valid the change is. Before opening, confirm there is
no other open PR for the same issue; if there is, explain how this PR differs.
### `### Contribution Checklist`
Check every item that applies. For the breaking-change item:
- Leave **"This is not a breaking change."** checked for the common case.
- If the change **is** breaking, add the `breaking change` label **or** put
`[BREAKING]` in the title prefix, before or after a language prefix such as
`Python:` or `.NET:` — workflows keep the label and the title prefix in sync
automatically (see `.github/workflows/label-title-prefix.yml` and
`.github/workflows/label-pr.yml`).
### Do not
- Do **not** add ad-hoc sections such as "Validation" or "Tests run"; CI/CD and
the checklist already cover validation status.
- Do **not** remove or reorder the template's headings.
### Creating the PR
Open new PRs as **drafts** until they are ready for review. Example:
```bash
gh pr create --repo microsoft/agent-framework --base main \
--head <your-fork-owner>:<branch> --draft \
--title "<concise title>" --body "<body following the template>"
```
## 2. Handling review comments
When a PR receives review comments, follow this sequence — **do not start editing
code before the user has reviewed the plan**:
1. **Review the comments.** Read every review comment and thread on the PR,
including inline code comments and general review summaries.
2. **Make a plan.** Produce a concrete plan describing how each comment will be
addressed (or why it should not be, with reasoning).
3. **Let the user review the plan.** Present the plan and wait for the user's
approval or adjustments before implementing anything.
4. **Implement.** Make the agreed changes.
5. **Reply to every comment.** Add a reply to **all** comments explaining how it
was addressed (or the agreed outcome) — leave none unanswered.
6. **Resolve resolved threads.** Mark a review thread as resolved only when the
comment has actually been addressed.
### Useful commands
List review comments and threads:
```bash
# Inline review comments
gh api repos/{owner}/{repo}/pulls/{pr}/comments
# Review threads with resolution state (GraphQL)
gh api graphql -f query='
query($owner:String!,$repo:String!,$pr:Int!){
repository(owner:$owner,name:$repo){
pullRequest(number:$pr){
reviewThreads(first:100){
nodes{ id isResolved comments(first:50){ nodes{ id body author{login} } } }
}
}
}
}' -F owner={owner} -F repo={repo} -F pr={pr}
```
Reply to an inline review comment:
```bash
gh api repos/{owner}/{repo}/pulls/{pr}/comments/{comment_id}/replies \
-f body="Addressed in <commit>: <explanation>"
```
Resolve a review thread (needs the thread node id from the GraphQL query above):
```bash
gh api graphql -f query='
mutation($threadId:ID!){
resolveReviewThread(input:{threadId:$threadId}){ thread{ isResolved } }
}' -F threadId={thread_id}
```
@@ -48,10 +48,6 @@ jobs:
filters: |
dotnet:
- 'dotnet/**'
- '!dotnet/AGENTS.md'
- '!dotnet/**/AGENTS.md'
- '!dotnet/.github/skills/*'
- '!dotnet/.github/skills/**'
cosmosdb:
- 'dotnet/src/Microsoft.Agents.AI.CosmosNoSql/**'
# The Foundry hosted-agent IT is costly (builds a container, pushes to ACR,
-4
View File
@@ -10,10 +10,6 @@ on:
branches: ["main", "feature*"]
paths:
- dotnet/**
- '!dotnet/AGENTS.md'
- '!dotnet/**/AGENTS.md'
- '!dotnet/.github/skills/*'
- '!dotnet/.github/skills/**'
- '.github/workflows/dotnet-format.yml'
concurrency:
+1 -19
View File
@@ -6,34 +6,16 @@
# https://github.com/actions/labeler
name: Label pull request
on:
pull_request_target:
types: [opened, synchronize, reopened, edited]
on: [pull_request_target]
jobs:
add_label:
runs-on: ubuntu-latest
permissions:
contents: read
issues: write
pull-requests: write
steps:
- uses: actions/labeler@f27b608878404679385c85cfa523b85ccb86e213 # v6
with:
repo-token: "${{ secrets.GH_ACTIONS_PR_WRITE }}"
- name: Checkout scripts
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
with:
sparse-checkout: .github/scripts
fetch-depth: 1
persist-credentials: false
- name: "PR: add breaking change label from title"
uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8
with:
github-token: ${{ secrets.GH_ACTIONS_PR_WRITE }}
script: |
const { syncBreakingChangeLabelFromTitle } = require('./.github/scripts/title_prefix.js');
await syncBreakingChangeLabelFromTitle({ github, context, core });
+50 -9
View File
@@ -15,17 +15,58 @@ jobs:
pull-requests: write
steps:
- name: Checkout scripts
uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6
with:
sparse-checkout: .github/scripts
fetch-depth: 1
persist-credentials: false
- uses: actions/github-script@ed597411d8f924073f98dfc5c65a23a2325f34cd # v8
name: "Issue/PR: update title"
with:
github-token: ${{ secrets.GITHUB_TOKEN }}
script: |
const { updateTitleForAddedLabel } = require('./.github/scripts/title_prefix.js');
await updateTitleForAddedLabel({ github, context, core });
let prefixLabels = {
"python": "Python",
".NET": ".NET"
};
function addTitlePrefix(title, prefix)
{
// Update the title based on the label and prefix
// Check if the title starts with the prefix (case-sensitive)
if (!title.startsWith(prefix + ": ")) {
// If not, check if the first word is the label (case-insensitive)
if (title.match(new RegExp(`^${prefix}`, 'i'))) {
// If yes, replace it with the prefix (case-sensitive)
title = title.replace(new RegExp(`^${prefix}`, 'i'), prefix);
} else {
// If not, prepend the prefix to the title
title = prefix + ": " + title;
}
}
return title;
}
labelAdded = context.payload.label.name
// Check if the issue or PR has the label
if (labelAdded in prefixLabels) {
let prefix = prefixLabels[labelAdded];
switch(context.eventName) {
case 'issues':
github.rest.issues.update({
issue_number: context.issue.number,
owner: context.repo.owner,
repo: context.repo.repo,
title: addTitlePrefix(context.payload.issue.title, prefix)
});
break
case 'pull_request_target':
github.rest.pulls.update({
pull_number: context.issue.number,
owner: context.repo.owner,
repo: context.repo.repo,
title: addTitlePrefix(context.payload.pull_request.title, prefix)
});
break
default:
core.setFailed('Unrecognited eventName: ' + context.eventName);
}
}
@@ -6,10 +6,6 @@ on:
branches: ["main"]
paths:
- "python/**"
- "!python/AGENTS.md"
- "!python/**/AGENTS.md"
- "!python/.github/skills/*"
- "!python/.github/skills/**"
env:
# Configure a constant location for the uv cache
-4
View File
@@ -31,10 +31,6 @@ jobs:
filters: |
python:
- 'python/**'
- '!python/AGENTS.md'
- '!python/**/AGENTS.md'
- '!python/.github/skills/*'
- '!python/.github/skills/**'
# run only if 'python' files were changed
- name: python tests
if: steps.filter.outputs.python == 'true'
-4
View File
@@ -49,10 +49,6 @@ jobs:
filters: |
python:
- 'python/**'
- '!python/AGENTS.md'
- '!python/**/AGENTS.md'
- '!python/.github/skills/*'
- '!python/.github/skills/**'
- '.github/actions/setup-local-mcp-server/**'
- '.github/workflows/python-merge-tests.yml'
- '.github/workflows/python-integration-tests.yml'
-4
View File
@@ -5,10 +5,6 @@ on:
branches: ["main", "feature*"]
paths:
- "python/**"
- "!python/AGENTS.md"
- "!python/**/AGENTS.md"
- "!python/.github/skills/*"
- "!python/.github/skills/**"
env:
# Configure a constant location for the uv cache
UV_CACHE_DIR: /tmp/.uv-cache
-1
View File
@@ -1 +0,0 @@
../../../.github/skills/pull-requests
-4
View File
@@ -10,10 +10,6 @@ See `./.github/skills/build-and-test/SKILL.md` for detailed instructions on buil
See `./.github/skills/project-structure/SKILL.md` for an overview of the project structure.
## Pull Requests
See `./.github/skills/pull-requests/SKILL.md` for guidance on writing PR descriptions and handling/resolving PR review comments.
### Core types
- `AIAgent`: The abstract base class that all agents derive from, providing common methods for interacting with an agent.
-1
View File
@@ -21,7 +21,6 @@
<PackageVersion Include="Aspire.Hosting.Azure.CognitiveServices" Version="$(AspireAppHostSdkVersion)" />
<PackageVersion Include="Aspire.Microsoft.Azure.Cosmos" Version="$(AspireAppHostSdkVersion)" />
<PackageVersion Include="CommunityToolkit.Aspire.OllamaSharp" Version="13.0.0" />
<PackageVersion Include="MessagePack" Version="3.1.7" /> <!-- Transitive dependency of Aspire pinned to newer version due to vulnerability in 2.5.192 -->
<!-- Azure.* -->
<PackageVersion Include="Azure.AI.AgentServer.Core" Version="1.0.0-beta.25" />
<PackageVersion Include="Azure.AI.AgentServer.Invocations" Version="1.0.0-beta.4" />
@@ -21,8 +21,6 @@ public sealed class PlanningOutputObserver : ConsoleObserver
private readonly string _planModeName;
private readonly string _executionModeName;
private readonly IReadOnlyDictionary<string, ConsoleColor>? _modeColors;
private string? _lastResponseId;
private string? _lastMessageId;
/// <summary>
/// Initializes a new instance of the <see cref="PlanningOutputObserver"/> class.
@@ -49,38 +47,17 @@ public sealed class PlanningOutputObserver : ConsoleObserver
}
/// <inheritdoc/>
public override async Task OnResponseUpdateAsync(IUXStateDriver ux, AgentResponseUpdate update, AIAgent agent, AgentSession session)
public override Task OnTextAsync(IUXStateDriver ux, string text, AIAgent agent, AgentSession session)
{
// We aren't in planning mode, so we can just stream the output directly.
if (!this.IsPlanningMode(ux.CurrentMode))
if (this.IsPlanningMode(ux.CurrentMode))
{
if (!string.IsNullOrWhiteSpace(update.Text))
{
await ux.WriteTextAsync(update.Text).ConfigureAwait(false);
}
return;
// Planning mode: collect text silently for JSON parsing after the stream.
this._textCollector.Append(text);
return Task.CompletedTask;
}
// We are still accumulating the same response/message.
if (this._lastResponseId == update.ResponseId && this._lastMessageId == update.MessageId)
{
this._textCollector.Append(update.Text);
return;
}
// New response/message, write the previous response/message and
// clear the text collector for the next JSON response/message.
string collectedText = this._textCollector.ToString();
if (!string.IsNullOrWhiteSpace(collectedText))
{
await ux.WriteTextAsync(collectedText).ConfigureAwait(false);
}
this._textCollector.Clear();
this._textCollector.Append(update.Text);
this._lastResponseId = update.ResponseId;
this._lastMessageId = update.MessageId;
// Execution mode: stream text directly to the console.
return ux.WriteTextAsync(text);
}
/// <inheritdoc/>
@@ -46,13 +46,6 @@ public class PlanningQuestion
/// Only used for clarification questions. Null when no predefined choices are offered.
/// </summary>
[JsonPropertyName("choices")]
[Description("""
For clarifications, this has a list of options that the user can choose from.
null for approvals.
Note: for clarifications, the user will always also be presented with a free form input option, so make sure that each choice provided here is a valid input for the next turn.
E.g. if the question is "Which stock are you referring to?" then valid choices might be ["AAPL", "MSFT", "GOOG"], and the user could also type their own answer.
Invalid choices would be ["Enter tickers directly", "Paste tickers"], since these conflict with the already existing freeform option, and don't directly provide valid inputs for the next turn.
""")]
[Description("For clarifications, this has a list of options that the user can choose from. null for approvals.")]
public List<string>? Choices { get; set; }
}
@@ -34,10 +34,10 @@ using var tracerProvider = HarnessTracing.CreateFileTracerProvider(TracingSource
var instructions =
"""
You are a data analyst assistant. You have access to a folder of data files via the file_access_* tools.
You are a data analyst assistant. You have access to a folder of data files via the FileAccess_* tools.
## Getting started
- Start by listing available files with file_access_list_files to see what data is available.
- Start by listing available files with FileAccess_ListFiles to see what data is available.
- Read the files to understand their structure and contents.
## Working with data
@@ -46,7 +46,7 @@ var instructions =
- When calculations are needed, work through them step by step and show your reasoning.
## Writing output
- When asked to produce output files (e.g., reports, summaries, filtered data), use file_access_save_file to write them.
- When asked to produce output files (e.g., reports, summaries, filtered data), use FileAccess_SaveFile to write them.
- Use appropriate file formats: CSV for tabular data, Markdown for reports.
- Confirm what you wrote and where.
@@ -8,7 +8,6 @@ using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
using Azure.Core;
using Microsoft.Agents.AI.CosmosNoSql;
using Microsoft.Azure.Cosmos;
using Microsoft.Extensions.AI;
using Microsoft.Shared.Diagnostics;
@@ -127,7 +126,6 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
Throw.IfNull(stateInitializer),
stateKey ?? this.GetType().Name);
this._cosmosClient = Throw.IfNull(cosmosClient);
CosmosOptionsHelper.EnsureApplicationName(this._cosmosClient, nameof(CosmosChatHistoryProvider));
this.DatabaseId = Throw.IfNullOrWhitespace(databaseId);
this.ContainerId = Throw.IfNullOrWhitespace(containerId);
this._container = this._cosmosClient.GetContainer(databaseId, containerId);
@@ -159,7 +157,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputRequestMessageFilter = null,
Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputResponseMessageFilter = null)
: this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString), CosmosOptionsHelper.CreateOptions(nameof(CosmosChatHistoryProvider))), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputRequestMessageFilter, storeInputResponseMessageFilter)
: this(new CosmosClient(Throw.IfNullOrWhitespace(connectionString)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputRequestMessageFilter, storeInputResponseMessageFilter)
{
}
@@ -187,7 +185,7 @@ public sealed class CosmosChatHistoryProvider : ChatHistoryProvider, IDisposable
Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? provideOutputMessageFilter = null,
Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputRequestMessageFilter = null,
Func<IEnumerable<ChatMessage>, IEnumerable<ChatMessage>>? storeInputResponseMessageFilter = null)
: this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential), CosmosOptionsHelper.CreateOptions(nameof(CosmosChatHistoryProvider))), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputRequestMessageFilter, storeInputResponseMessageFilter)
: this(new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential)), databaseId, containerId, stateInitializer, ownsClient: true, stateKey, provideOutputMessageFilter, storeInputRequestMessageFilter, storeInputResponseMessageFilter)
{
}
@@ -7,7 +7,6 @@ using System.Linq;
using System.Text.Json;
using System.Threading.Tasks;
using Azure.Core;
using Microsoft.Agents.AI.CosmosNoSql;
using Microsoft.Azure.Cosmos;
using Microsoft.Shared.Diagnostics;
using Newtonsoft.Json;
@@ -38,7 +37,7 @@ public class CosmosCheckpointStore<T> : JsonCheckpointStore, IDisposable
/// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
public CosmosCheckpointStore(string connectionString, string databaseId, string containerId)
{
var cosmosClientOptions = CosmosOptionsHelper.CreateOptions(nameof(CosmosCheckpointStore));
var cosmosClientOptions = new CosmosClientOptions();
this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(connectionString), cosmosClientOptions);
this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId));
@@ -56,10 +55,12 @@ public class CosmosCheckpointStore<T> : JsonCheckpointStore, IDisposable
/// <exception cref="ArgumentException">Thrown when any string parameter is null or whitespace.</exception>
public CosmosCheckpointStore(string accountEndpoint, TokenCredential tokenCredential, string databaseId, string containerId)
{
var cosmosClientOptions = CosmosOptionsHelper.CreateOptions(nameof(CosmosCheckpointStore));
cosmosClientOptions.SerializerOptions = new CosmosSerializationOptions
var cosmosClientOptions = new CosmosClientOptions
{
PropertyNamingPolicy = CosmosPropertyNamingPolicy.CamelCase
SerializerOptions = new CosmosSerializationOptions
{
PropertyNamingPolicy = CosmosPropertyNamingPolicy.CamelCase
}
};
this._cosmosClient = new CosmosClient(Throw.IfNullOrWhitespace(accountEndpoint), Throw.IfNull(tokenCredential), cosmosClientOptions);
@@ -78,7 +79,6 @@ public class CosmosCheckpointStore<T> : JsonCheckpointStore, IDisposable
public CosmosCheckpointStore(CosmosClient cosmosClient, string databaseId, string containerId)
{
this._cosmosClient = Throw.IfNull(cosmosClient);
CosmosOptionsHelper.EnsureApplicationName(this._cosmosClient, nameof(CosmosCheckpointStore));
this._container = this._cosmosClient.GetContainer(Throw.IfNullOrWhitespace(databaseId), Throw.IfNullOrWhitespace(containerId));
this._ownsClient = false;
@@ -1,80 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Reflection;
using Microsoft.Azure.Cosmos;
namespace Microsoft.Agents.AI.CosmosNoSql;
/// <summary>
/// Provides shared Cosmos DB client configuration for Agent Framework Cosmos NoSQL integrations.
/// Ensures all internally-created <see cref="CosmosClient"/> instances carry a consistent
/// <see cref="CosmosClientOptions.ApplicationName"/> for telemetry and diagnostics.
/// </summary>
internal static class CosmosOptionsHelper
{
/// <summary>
/// Maximum length allowed by the Cosmos DB .NET SDK for <see cref="CosmosClientOptions.ApplicationName"/>.
/// </summary>
private const int MaxApplicationNameLength = 64;
private static readonly string s_version = GetVersion();
/// <summary>
/// Creates a <see cref="CosmosClientOptions"/> instance pre-configured with the
/// Agent Framework application name for User-Agent identification.
/// </summary>
/// <param name="component">The fully-qualified component class name (e.g. "CosmosChatHistoryProvider").</param>
/// <returns>A new <see cref="CosmosClientOptions"/> with <see cref="CosmosClientOptions.ApplicationName"/> set.</returns>
public static CosmosClientOptions CreateOptions(string component)
{
return new CosmosClientOptions
{
ApplicationName = BuildApplicationName(component)
};
}
/// <summary>
/// Ensures the given <see cref="CosmosClient"/> has an <see cref="CosmosClientOptions.ApplicationName"/> set.
/// If the client already has a non-empty ApplicationName, it is not overridden.
/// </summary>
/// <param name="cosmosClient">The client to apply the application name to.</param>
/// <param name="component">The fully-qualified component class name (e.g. "CosmosChatHistoryProvider").</param>
public static void EnsureApplicationName(CosmosClient cosmosClient, string component)
{
if (string.IsNullOrWhiteSpace(cosmosClient.ClientOptions.ApplicationName))
{
cosmosClient.ClientOptions.ApplicationName = BuildApplicationName(component);
}
}
private static string BuildApplicationName(string component)
{
var applicationName = $"Microsoft.Agents.AI.CosmosNoSql.{component}/{s_version}";
if (applicationName.Length > MaxApplicationNameLength)
{
applicationName = applicationName.Substring(0, MaxApplicationNameLength);
}
return applicationName;
}
private static string GetVersion()
{
if (typeof(CosmosOptionsHelper).Assembly.GetCustomAttribute<AssemblyInformationalVersionAttribute>()?.InformationalVersion is string version)
{
int pos = version.IndexOf('+', System.StringComparison.Ordinal);
if (pos >= 0)
{
version = version.Substring(0, pos);
}
if (version.Length > 0)
{
return version;
}
}
return "unknown";
}
}
@@ -35,7 +35,7 @@ internal sealed class ClientHeadersAgent : DelegatingAIAgent
}
/// <inheritdoc/>
protected override async Task<AgentResponse> RunCoreAsync(
protected override Task<AgentResponse> RunCoreAsync(
IEnumerable<ChatMessage> messages,
AgentSession? session = null,
AgentRunOptions? options = null,
@@ -44,15 +44,13 @@ internal sealed class ClientHeadersAgent : DelegatingAIAgent
var snapshot = TrySnapshot(options);
if (snapshot is not null)
{
// This method is async, so the runtime restores the caller's ExecutionContext (and
// therefore the previous ClientHeadersScope.Current value) when the returned task
// completes. Awaiting the inner call is what establishes that async-method boundary,
// so the per-run scope set here cannot carry into a later run on the same async flow.
// See ClientHeadersScope remarks. The streaming path relies on the same behavior.
// AsyncLocal mutations made inside an awaited async method do not leak back to the
// caller after the method returns, so we do not need an explicit restore step here.
// See ClientHeadersScope remarks.
ClientHeadersScope.Current = snapshot;
}
return await this.InnerAgent.RunAsync(messages, session, options, cancellationToken).ConfigureAwait(false);
return this.InnerAgent.RunAsync(messages, session, options, cancellationToken);
}
/// <inheritdoc/>
@@ -31,12 +31,11 @@ namespace Microsoft.Agents.AI;
/// <para>
/// This provider exposes the following tools to the agent:
/// <list type="bullet">
/// <item><description><c>file_access_save_file</c> — Save a file with the given name and content.</description></item>
/// <item><description><c>file_access_read_file</c> — Read the content of a file by name.</description></item>
/// <item><description><c>file_access_delete_file</c> — Delete a file by name.</description></item>
/// <item><description><c>file_access_list_files</c> — List the direct child file names in a directory.</description></item>
/// <item><description><c>file_access_list_subdirectories</c> — List the direct child subdirectory names in a directory.</description></item>
/// <item><description><c>file_access_search_files</c> — Recursively search file contents using a regular expression pattern.</description></item>
/// <item><description><c>SaveFile</c> — Save a file with the given name and content.</description></item>
/// <item><description><c>ReadFile</c> — Read the content of a file by name.</description></item>
/// <item><description><c>DeleteFile</c> — Delete a file by name.</description></item>
/// <item><description><c>ListFiles</c> — List all file names.</description></item>
/// <item><description><c>SearchFiles</c> — Search file contents using a regular expression pattern.</description></item>
/// </list>
/// </para>
/// </remarks>
@@ -46,13 +45,11 @@ public sealed class FileAccessProvider : AIContextProvider
private const string DefaultInstructions =
"""
## File Access
You have access to a shared file storage area via the `file_access_*` tools for reading, writing, and managing files.
You have access to a shared file storage area via the `FileAccess_*` tools for reading, writing, and managing files.
These files persist beyond the current session and may be shared across sessions or agents.
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.
- Never delete or overwrite existing files unless the user has explicitly asked you to do so.
- 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.
""";
private readonly AgentFileStore _fileStore;
@@ -140,56 +137,30 @@ public sealed class FileAccessProvider : AIContextProvider
}
/// <summary>
/// List the direct child file names of a directory. Omit <paramref name="directory"/> (or pass an empty string)
/// to list the store root. To enumerate files in a subdirectory, pass its relative path.
/// List all file names.
/// </summary>
/// <param name="directory">The relative directory path to list. Omit or pass an empty string to list the store root.</param>
/// <param name="cancellationToken">A token to cancel the operation.</param>
/// <returns>A list of file names.</returns>
[Description("List the direct child file names of a directory. Omit the directory (or pass an empty string) to list the root. To enumerate files in a subdirectory, pass its relative path, for example \"reports\" or \"reports/2024\".")]
private async Task<List<string>> ListFilesAsync(string? directory = null, CancellationToken cancellationToken = default)
[Description("List all file names.")]
private async Task<List<string>> ListFilesAsync(CancellationToken cancellationToken = default)
{
string target = string.IsNullOrWhiteSpace(directory) ? string.Empty : directory;
IReadOnlyList<string> fileNames = await this._fileStore.ListFilesAsync(target, cancellationToken).ConfigureAwait(false);
IReadOnlyList<string> fileNames = await this._fileStore.ListFilesAsync(string.Empty, cancellationToken).ConfigureAwait(false);
return new List<string>(fileNames);
}
/// <summary>
/// List the direct child subdirectory names of a directory. Omit <paramref name="directory"/> (or pass an empty string)
/// to list the store root. To enumerate subdirectories of a subdirectory, pass its relative path.
/// </summary>
/// <param name="directory">The relative directory path to list. Omit or pass an empty string to list the store root.</param>
/// <param name="cancellationToken">A token to cancel the operation.</param>
/// <returns>A list of subdirectory names.</returns>
[Description("List the direct child subdirectory names of a directory. Omit the 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.")]
private async Task<List<string>> ListSubdirectoriesAsync(string? directory = null, CancellationToken cancellationToken = default)
{
string target = string.IsNullOrWhiteSpace(directory) ? string.Empty : directory;
IReadOnlyList<string> directoryNames = await this._fileStore.ListDirectoriesAsync(target, cancellationToken).ConfigureAwait(false);
return new List<string>(directoryNames);
}
/// <summary>
/// Search the contents of all files in the store (recursively) using a regular expression pattern (case-insensitive).
/// Search file contents using a regular expression pattern (case-insensitive).
/// Optionally filter which files to search using a glob pattern.
/// </summary>
/// <param name="regexPattern">A regular expression pattern to match against file contents (case-insensitive).</param>
/// <param name="filePattern">An optional glob pattern to filter which files to search, matched against each file's path relative to the store root. Use <c>**</c> to match across subdirectories (e.g., "**/*.md"). Leave empty or omit to search all files.</param>
/// <param name="filePattern">An optional glob pattern to filter which files to search (e.g., "*.md", "research*"). Leave empty or omit to search all files.</param>
/// <param name="cancellationToken">A token to cancel the operation.</param>
/// <returns>A list of search results whose file names are paths relative to the store root.</returns>
[Description(
"""
Search the contents of all files in the store (recursively, across all subdirectories) using a regular expression pattern (case-insensitive).
Optionally filter which files to search using a glob pattern matched against each file's path relative to the store root:
- '*' matches within a single path segment
- '**' matches across subdirectories, so use \"**/*.md\" to match markdown files at any depth, or \"reports/**\" to restrict the search to the 'reports' subtree.
Returns matching results whose file names are paths relative to the store root (usable with file_access_read_file), along with snippets and matching lines with line numbers.
""")]
/// <returns>A list of search results with matching file names, snippets, and matching lines.</returns>
[Description("Search file contents using a regular expression pattern (case-insensitive). Optionally filter which files to search using a glob pattern (e.g., \"*.md\", \"research*\"). Returns matching file names, snippets, and matching lines with line numbers.")]
private async Task<List<FileSearchResult>> SearchFilesAsync(string regexPattern, string? filePattern = null, CancellationToken cancellationToken = default)
{
string? pattern = string.IsNullOrWhiteSpace(filePattern) ? null : filePattern;
IReadOnlyList<FileSearchResult> results = await this._fileStore.SearchFilesAsync(string.Empty, regexPattern, pattern, recursive: true, cancellationToken).ConfigureAwait(false);
IReadOnlyList<FileSearchResult> results = await this._fileStore.SearchFilesAsync(string.Empty, regexPattern, pattern, cancellationToken).ConfigureAwait(false);
return new List<FileSearchResult>(results);
}
@@ -199,12 +170,11 @@ public sealed class FileAccessProvider : AIContextProvider
return
[
AIFunctionFactory.Create(this.SaveFileAsync, new AIFunctionFactoryOptions { Name = "file_access_save_file", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.ReadFileAsync, new AIFunctionFactoryOptions { Name = "file_access_read_file", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.DeleteFileAsync, new AIFunctionFactoryOptions { Name = "file_access_delete_file", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.ListFilesAsync, new AIFunctionFactoryOptions { Name = "file_access_list_files", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.ListSubdirectoriesAsync, new AIFunctionFactoryOptions { Name = "file_access_list_subdirectories", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.SearchFilesAsync, new AIFunctionFactoryOptions { Name = "file_access_search_files", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.SaveFileAsync, new AIFunctionFactoryOptions { Name = "FileAccess_SaveFile", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.ReadFileAsync, new AIFunctionFactoryOptions { Name = "FileAccess_ReadFile", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.DeleteFileAsync, new AIFunctionFactoryOptions { Name = "FileAccess_DeleteFile", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.ListFilesAsync, new AIFunctionFactoryOptions { Name = "FileAccess_ListFiles", SerializerOptions = serializerOptions }),
AIFunctionFactory.Create(this.SearchFilesAsync, new AIFunctionFactoryOptions { Name = "FileAccess_SearchFiles", SerializerOptions = serializerOptions }),
];
}
}
@@ -296,7 +296,7 @@ public sealed class FileMemoryProvider : AIContextProvider, IDisposable
{
FileMemoryState state = this._sessionState.GetOrInitializeState(AIAgent.CurrentRunContext?.Session);
string? pattern = string.IsNullOrWhiteSpace(filePattern) ? null : filePattern;
IReadOnlyList<FileSearchResult> results = await this._fileStore.SearchFilesAsync(state.WorkingFolder, regexPattern, pattern, recursive: false, cancellationToken).ConfigureAwait(false);
IReadOnlyList<FileSearchResult> results = await this._fileStore.SearchFilesAsync(state.WorkingFolder, regexPattern, pattern, cancellationToken).ConfigureAwait(false);
// Filter out internal files (description sidecars and memory index) so they stay hidden.
var filtered = new List<FileSearchResult>(results.Count);
@@ -58,14 +58,6 @@ public abstract class AgentFileStore
/// <returns>A list of file names in the specified directory (direct children only).</returns>
public abstract Task<IReadOnlyList<string>> ListFilesAsync(string directory, CancellationToken cancellationToken = default);
/// <summary>
/// Lists the direct child subdirectories of a directory.
/// </summary>
/// <param name="directory">The relative path of the directory to list. Use an empty string for the root.</param>
/// <param name="cancellationToken">A token to cancel the operation.</param>
/// <returns>A list of subdirectory names in the specified directory (direct children only).</returns>
public abstract Task<IReadOnlyList<string>> ListDirectoriesAsync(string directory, CancellationToken cancellationToken = default);
/// <summary>
/// Checks whether a file exists.
/// </summary>
@@ -84,20 +76,12 @@ public abstract class AgentFileStore
/// </param>
/// <param name="filePattern">
/// An optional glob pattern to filter which files are searched (e.g., <c>"*.md"</c>, <c>"research*"</c>).
/// When <see langword="null"/>, all files are searched.
/// Uses standard glob syntax from <see cref="Matcher"/>, matched against each file's path relative to
/// <paramref name="directory"/>. Use <c>**</c> to match across subdirectories (e.g., <c>"**/*.md"</c>).
/// </param>
/// <param name="recursive">
/// When <see langword="true"/>, all descendant files of <paramref name="directory"/> are searched.
/// When <see langword="false"/> (default), only the direct children of <paramref name="directory"/> are searched.
/// When <see langword="null"/>, all files in the directory are searched.
/// Uses standard glob syntax from <see cref="Matcher"/>.
/// </param>
/// <param name="cancellationToken">A token to cancel the operation.</param>
/// <returns>
/// A list of search results. Each result's <see cref="FileSearchResult.FileName"/> is the matching file's
/// path relative to <paramref name="directory"/>.
/// </returns>
public abstract Task<IReadOnlyList<FileSearchResult>> SearchFilesAsync(string directory, string regexPattern, string? filePattern = null, bool recursive = false, CancellationToken cancellationToken = default);
/// <returns>A list of search results with matching file names, snippets, and matching lines.</returns>
public abstract Task<IReadOnlyList<FileSearchResult>> SearchFilesAsync(string directory, string regexPattern, string? filePattern = null, CancellationToken cancellationToken = default);
/// <summary>
/// Ensures a directory exists, creating it if necessary.
@@ -142,7 +142,6 @@ public sealed class FileSystemAgentFileStore : AgentFileStore
string directory,
string regexPattern,
string? filePattern = null,
bool recursive = false,
CancellationToken cancellationToken = default)
{
string fullDir = this.ResolveSafeDirectoryPath(directory);
@@ -157,13 +156,22 @@ public sealed class FileSystemAgentFileStore : AgentFileStore
Matcher? matcher = filePattern is not null ? StorePaths.CreateGlobMatcher(filePattern) : null;
var results = new List<FileSearchResult>();
foreach (string filePath in EnumerateFiles(fullDir, recursive))
foreach (string filePath in Directory.GetFiles(fullDir))
{
// The file path relative to the search directory, using forward slashes.
string relativeName = GetRelativeStorePath(fullDir, filePath);
// Skip files that are symlinks/reparse points to prevent reading outside the root.
if ((File.GetAttributes(filePath) & FileAttributes.ReparsePoint) != 0)
{
continue;
}
// Apply the optional glob filter on the relative path.
if (!StorePaths.MatchesGlob(relativeName, matcher))
string? fileName = Path.GetFileName(filePath);
if (fileName is null)
{
continue;
}
// Apply the optional glob filter on the file name.
if (!StorePaths.MatchesGlob(fileName, matcher))
{
continue;
}
@@ -210,7 +218,7 @@ public sealed class FileSystemAgentFileStore : AgentFileStore
{
results.Add(new FileSearchResult
{
FileName = relativeName,
FileName = fileName,
Snippet = firstSnippet!,
MatchingLines = matchingLines,
});
@@ -220,76 +228,6 @@ public sealed class FileSystemAgentFileStore : AgentFileStore
return results;
}
/// <inheritdoc />
public override Task<IReadOnlyList<string>> ListDirectoriesAsync(string directory, CancellationToken cancellationToken = default)
{
string fullDir = this.ResolveSafeDirectoryPath(directory);
if (!Directory.Exists(fullDir))
{
return Task.FromResult<IReadOnlyList<string>>([]);
}
var directories = Directory.GetDirectories(fullDir)
.Where(d => (File.GetAttributes(d) & FileAttributes.ReparsePoint) == 0)
.Select(Path.GetFileName)
.Where(name => name is not null)
.ToList();
return Task.FromResult<IReadOnlyList<string>>(directories!);
}
/// <summary>
/// Enumerates the files directly under <paramref name="directory"/> (or all descendant files when
/// <paramref name="recursive"/> is <see langword="true"/>), skipping symlinks/reparse points for both
/// files and directories to prevent reading outside the root.
/// </summary>
private static IEnumerable<string> EnumerateFiles(string directory, bool recursive)
{
foreach (string filePath in Directory.EnumerateFiles(directory))
{
// Skip files that are symlinks/reparse points.
if ((File.GetAttributes(filePath) & FileAttributes.ReparsePoint) != 0)
{
continue;
}
yield return filePath;
}
if (!recursive)
{
yield break;
}
foreach (string subDir in Directory.EnumerateDirectories(directory))
{
// Skip symlinked/reparse-point directories so recursion cannot escape the root.
if ((File.GetAttributes(subDir) & FileAttributes.ReparsePoint) != 0)
{
continue;
}
foreach (string filePath in EnumerateFiles(subDir, recursive: true))
{
yield return filePath;
}
}
}
/// <summary>
/// Returns the path of <paramref name="filePath"/> relative to <paramref name="baseDirectory"/>,
/// normalized to forward-slash separators. Assumes <paramref name="filePath"/> resides under
/// <paramref name="baseDirectory"/> (as produced by <see cref="EnumerateFiles"/>).
/// </summary>
private static string GetRelativeStorePath(string baseDirectory, string filePath)
{
string baseTrimmed = baseDirectory.TrimEnd(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar);
string relative = filePath.Substring(baseTrimmed.Length)
.TrimStart(Path.DirectorySeparatorChar, Path.AltDirectorySeparatorChar);
return relative.Replace(Path.DirectorySeparatorChar, '/').Replace(Path.AltDirectorySeparatorChar, '/');
}
/// <inheritdoc />
public override Task CreateDirectoryAsync(string path, CancellationToken cancellationToken = default)
{
@@ -66,43 +66,6 @@ public sealed class InMemoryAgentFileStore : AgentFileStore
return Task.FromResult<IReadOnlyList<string>>(files);
}
/// <inheritdoc />
public override Task<IReadOnlyList<string>> ListDirectoriesAsync(string directory, CancellationToken cancellationToken = default)
{
string prefix = StorePaths.NormalizeRelativePath(directory, isDirectory: true);
if (prefix.Length > 0 && !prefix.EndsWith("/", StringComparison.Ordinal))
{
prefix += "/";
}
// A subdirectory is the first path segment of any key whose remainder (after the prefix)
// still contains a separator. Collect distinct first segments, preserving original casing.
var directories = new List<string>();
var seen = new HashSet<string>(StringComparer.OrdinalIgnoreCase);
foreach (string key in this._files.Keys)
{
if (!key.StartsWith(prefix, StringComparison.OrdinalIgnoreCase))
{
continue;
}
string remainder = key.Substring(prefix.Length);
int separatorIndex = remainder.IndexOf("/", StringComparison.Ordinal);
if (separatorIndex <= 0)
{
continue;
}
string segment = remainder.Substring(0, separatorIndex);
if (seen.Add(segment))
{
directories.Add(segment);
}
}
return Task.FromResult<IReadOnlyList<string>>(directories);
}
/// <inheritdoc />
public override Task<bool> FileExistsAsync(string path, CancellationToken cancellationToken = default)
{
@@ -111,7 +74,7 @@ public sealed class InMemoryAgentFileStore : AgentFileStore
}
/// <inheritdoc />
public override Task<IReadOnlyList<FileSearchResult>> SearchFilesAsync(string directory, string regexPattern, string? filePattern = null, bool recursive = false, CancellationToken cancellationToken = default)
public override Task<IReadOnlyList<FileSearchResult>> SearchFilesAsync(string directory, string regexPattern, string? filePattern = null, CancellationToken cancellationToken = default)
{
// Normalize the directory prefix for path matching.
string prefix = StorePaths.NormalizeRelativePath(directory, isDirectory: true);
@@ -133,16 +96,14 @@ public sealed class InMemoryAgentFileStore : AgentFileStore
continue;
}
// The file path relative to the search directory.
// Exclude files in subdirectories (direct children only).
string relativeName = kvp.Key.Substring(prefix.Length);
// When not recursive, exclude files in subdirectories (direct children only).
if (!recursive && relativeName.IndexOf("/", StringComparison.Ordinal) >= 0)
if (relativeName.IndexOf("/", StringComparison.Ordinal) >= 0)
{
continue;
}
// Apply the optional glob filter on the relative path.
// Apply the optional glob filter on the file name.
if (!StorePaths.MatchesGlob(relativeName, matcher))
{
continue;
@@ -1,75 +0,0 @@
// Copyright (c) Microsoft. All rights reserved.
using Microsoft.Agents.AI.CosmosNoSql;
using Microsoft.Azure.Cosmos;
namespace Microsoft.Agents.AI.CosmosNoSql.UnitTests;
public sealed class CosmosOptionsHelperTests
{
[Fact]
public void CreateOptions_SetsApplicationName_WithComponentAndVersion()
{
// Act
var options = CosmosOptionsHelper.CreateOptions("CosmosChatHistoryProvider");
// Assert
Assert.NotNull(options.ApplicationName);
Assert.StartsWith("Microsoft.Agents.AI.CosmosNoSql.CosmosChatHistoryProvider/", options.ApplicationName);
}
[Fact]
public void CreateOptions_DifferentComponents_ProduceDifferentNames()
{
// Act
var chatOptions = CosmosOptionsHelper.CreateOptions("CosmosChatHistoryProvider");
var checkpointOptions = CosmosOptionsHelper.CreateOptions("CosmosCheckpointStore");
// Assert
Assert.NotEqual(chatOptions.ApplicationName, checkpointOptions.ApplicationName);
Assert.Contains("CosmosChatHistoryProvider", chatOptions.ApplicationName);
Assert.Contains("CosmosCheckpointStore", checkpointOptions.ApplicationName);
}
[Fact]
public void CreateOptions_ApplicationName_DoesNotExceedMaxLength()
{
// Use a deliberately long component name to trigger truncation
var longComponent = new string('X', 100);
// Act
var options = CosmosOptionsHelper.CreateOptions(longComponent);
// Assert
Assert.True(options.ApplicationName!.Length <= 64,
$"ApplicationName length {options.ApplicationName.Length} exceeds max 64");
}
[Fact]
public void EnsureApplicationName_SetsName_WhenClientHasNone()
{
// Arrange
var clientOptions = new CosmosClientOptions();
Assert.Null(clientOptions.ApplicationName);
// Act
var options = CosmosOptionsHelper.CreateOptions("CosmosChatHistoryProvider");
// Assert - verify the returned options have ApplicationName set
Assert.NotNull(options.ApplicationName);
Assert.NotEmpty(options.ApplicationName);
}
[Fact]
public void CreateOptions_ApplicationName_ContainsVersion()
{
// Act
var options = CosmosOptionsHelper.CreateOptions("CosmosChatHistoryProvider");
// Assert - should contain a "/" followed by version info
Assert.Contains("/", options.ApplicationName);
var parts = options.ApplicationName!.Split('/');
Assert.Equal(2, parts.Length);
Assert.False(string.IsNullOrWhiteSpace(parts[1]), "Version portion should not be empty");
}
}
@@ -4,7 +4,6 @@ using System;
using System.ClientModel;
using System.ClientModel.Primitives;
using System.Collections.Generic;
using System.Linq;
using System.Net;
using System.Net.Http;
using System.Reflection;
@@ -562,82 +561,6 @@ public sealed class ClientHeadersExtensionsTests
Assert.Equal(1, EntriesCount(policies!));
}
// -------------------------------------------------------------------------------------------
// 21. Non-streaming hardening: a non-streaming run must restore the ambient ClientHeadersScope
// on return so a previous run's x-client-* headers do not carry into a later headerless run
// on the same async flow. (Streaming already restores naturally via its async iterator.)
// -------------------------------------------------------------------------------------------
[Fact]
public async Task NonStreaming_DoesNotCarryClientHeadersToSubsequentRunAsync()
{
// Arrange: a probe inner agent records ClientHeadersScope.Current observed at each run.
var observed = new List<IReadOnlyDictionary<string, string>?>();
var inner = new ProbeAgent(_ =>
{
observed.Add(ClientHeadersScope.Current);
return Task.CompletedTask;
});
var agent = new ClientHeadersAgent(inner);
// Act: run 1 supplies a client header; run 2 supplies fresh, empty ChatOptions (no headers).
var run1Options = new ChatOptions();
run1Options.WithClientHeader("x-client-end-user-id", "alice");
await agent.RunAsync(messages: [], options: new ChatClientAgentRunOptions(run1Options));
// The scope must not carry back into the caller's flow after run 1 returns.
Assert.Null(ClientHeadersScope.Current);
var run2Options = new ChatOptions();
await agent.RunAsync(messages: [], options: new ChatClientAgentRunOptions(run2Options));
// Assert: run 1 observed "alice"; run 2 observed no headers (did not inherit run 1's value).
Assert.Equal(2, observed.Count);
Assert.NotNull(observed[0]);
Assert.Equal("alice", observed[0]!["x-client-end-user-id"]);
Assert.Null(observed[1]);
Assert.Null(ClientHeadersScope.Current);
}
// -------------------------------------------------------------------------------------------
// 22. End-to-end non-streaming: a second headerless run on the same async flow must not carry
// the first run's x-client-end-user-id onto the wire.
// -------------------------------------------------------------------------------------------
[Fact]
public async Task EndToEnd_NonStreaming_SecondRunDoesNotInheritHeaderOnWireAsync()
{
// Arrange: a real OpenAI ResponsesClient pointed at a recording handler.
using var handler = new RecordingHandler(MinimalResponseJson());
#pragma warning disable CA5399
using var http = new HttpClient(handler);
#pragma warning restore CA5399
var openAIOptions = new OpenAIClientOptions { Transport = new HttpClientPipelineTransport(http) };
var openAIClient = new OpenAIClient(new ApiKeyCredential("fake"), openAIOptions);
IChatClient chatClient = openAIClient.GetResponsesClient().AsIChatClient();
AIAgent agent = new ChatClientAgent(chatClient).AsBuilder().UseClientHeaders().Build();
// Act: run 1 carries x-client-end-user-id; run 2 supplies fresh options with no client headers.
var run1 = new ChatClientAgentRunOptions(new ChatOptions());
run1.ChatOptions!.WithClientHeader("x-client-end-user-id", "alice");
await agent.RunAsync("hi", options: run1);
var afterRun1 = handler.Requests.Count;
var run2 = new ChatClientAgentRunOptions(new ChatOptions());
await agent.RunAsync("hi", options: run2);
// Assert: run 1 stamped the header; none of the requests issued by run 2 carry it.
// (Assert per-run rather than on an exact total count, which would be brittle to
// any extra/internal SDK requests.)
Assert.True(afterRun1 > 0);
Assert.Equal("alice", handler.Requests[0].Headers["x-client-end-user-id"]);
var run2Requests = handler.Requests.Skip(afterRun1).ToList();
Assert.NotEmpty(run2Requests);
Assert.All(run2Requests, r => Assert.False(r.Headers.ContainsKey("x-client-end-user-id")));
}
// -------------------------------------------------------------------------------------------
// Helpers
// -------------------------------------------------------------------------------------------
@@ -40,8 +40,8 @@ public class FileAccessProviderTests
// Arrange
var tools = await CreateToolsAsync();
// Assert — 6 tools: SaveFile, ReadFile, DeleteFile, ListFiles, ListSubdirectories, SearchFiles
Assert.Equal(6, tools.Count());
// Assert — 5 tools: SaveFile, ReadFile, DeleteFile, ListFiles, SearchFiles
Assert.Equal(5, tools.Count());
}
[Fact]
@@ -61,7 +61,7 @@ public class FileAccessProviderTests
// Assert
Assert.NotNull(result.Instructions);
Assert.Contains("File Access", result.Instructions);
Assert.Contains("file_access_", result.Instructions);
Assert.Contains("FileAccess_", result.Instructions);
Assert.Contains("persist beyond the current session", result.Instructions);
}
@@ -108,7 +108,7 @@ public class FileAccessProviderTests
// Arrange
var store = new InMemoryAgentFileStore();
var tools = await CreateToolsAsync(store);
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act
await InvokeToolAsync(saveFile, new AIFunctionArguments
@@ -128,7 +128,7 @@ public class FileAccessProviderTests
// Arrange — FileAccessProvider should never create description sidecar files.
var store = new InMemoryAgentFileStore();
var tools = await CreateToolsAsync(store);
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act
await InvokeToolAsync(saveFile, new AIFunctionArguments
@@ -148,7 +148,7 @@ public class FileAccessProviderTests
// Arrange
var store = new InMemoryAgentFileStore();
var tools = await CreateToolsAsync(store);
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
await InvokeToolAsync(saveFile, new AIFunctionArguments
{
@@ -175,7 +175,7 @@ public class FileAccessProviderTests
// Arrange
var store = new InMemoryAgentFileStore();
var tools = await CreateToolsAsync(store);
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
await InvokeToolAsync(saveFile, new AIFunctionArguments
{
@@ -200,7 +200,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act
var result = await InvokeToolAsync(saveFile, new AIFunctionArguments
@@ -225,7 +225,7 @@ public class FileAccessProviderTests
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Stored content");
var tools = await CreateToolsAsync(store);
var readFile = GetTool(tools, "file_access_read_file");
var readFile = GetTool(tools, "FileAccess_ReadFile");
// Act
var result = await InvokeToolAsync(readFile, new AIFunctionArguments
@@ -243,7 +243,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var readFile = GetTool(tools, "file_access_read_file");
var readFile = GetTool(tools, "FileAccess_ReadFile");
// Act
var result = await InvokeToolAsync(readFile, new AIFunctionArguments
@@ -267,7 +267,7 @@ public class FileAccessProviderTests
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Content");
var tools = await CreateToolsAsync(store);
var deleteFile = GetTool(tools, "file_access_delete_file");
var deleteFile = GetTool(tools, "FileAccess_DeleteFile");
// Act
var result = await InvokeToolAsync(deleteFile, new AIFunctionArguments
@@ -286,7 +286,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var deleteFile = GetTool(tools, "file_access_delete_file");
var deleteFile = GetTool(tools, "FileAccess_DeleteFile");
// Act
var result = await InvokeToolAsync(deleteFile, new AIFunctionArguments
@@ -311,7 +311,7 @@ public class FileAccessProviderTests
await store.WriteFileAsync("notes.md", "Content");
await store.WriteFileAsync("data.txt", "Data");
var tools = await CreateToolsAsync(store);
var listFiles = GetTool(tools, "file_access_list_files");
var listFiles = GetTool(tools, "FileAccess_ListFiles");
// Act
var result = await InvokeToolAsync(listFiles, new AIFunctionArguments());
@@ -331,7 +331,7 @@ public class FileAccessProviderTests
await store.WriteFileAsync("notes.md", "Content");
await store.WriteFileAsync("notes_description.md", "Description");
var tools = await CreateToolsAsync(store);
var listFiles = GetTool(tools, "file_access_list_files");
var listFiles = GetTool(tools, "FileAccess_ListFiles");
// Act
var result = await InvokeToolAsync(listFiles, new AIFunctionArguments());
@@ -346,7 +346,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var listFiles = GetTool(tools, "file_access_list_files");
var listFiles = GetTool(tools, "FileAccess_ListFiles");
// Act
var result = await InvokeToolAsync(listFiles, new AIFunctionArguments());
@@ -356,98 +356,6 @@ public class FileAccessProviderTests
Assert.Empty(entries);
}
[Fact]
public async Task ListFiles_WithDirectory_ListsSubdirectoryChildrenAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("root.txt", "Root");
await store.WriteFileAsync("reports/2024/q1.md", "Q1");
await store.WriteFileAsync("reports/2024/q2.md", "Q2");
var tools = await CreateToolsAsync(store);
var listFiles = GetTool(tools, "file_access_list_files");
// Act
var result = await InvokeToolAsync(listFiles, new AIFunctionArguments
{
["directory"] = "reports/2024",
});
// Assert — only the direct children of reports/2024 are returned (by their names)
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
Assert.Equal(2, entries.Count);
Assert.Contains(entries, e => e.GetString() == "q1.md");
Assert.Contains(entries, e => e.GetString() == "q2.md");
}
#endregion
#region ListSubdirectories Tests
[Fact]
public async Task ListSubdirectories_ReturnsDirectChildDirectoriesAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("root.txt", "Root");
await store.WriteFileAsync("reports/q1.md", "Q1");
await store.WriteFileAsync("reports/2024/q2.md", "Q2");
await store.WriteFileAsync("data/raw.csv", "x");
var tools = await CreateToolsAsync(store);
var listSubdirectories = GetTool(tools, "file_access_list_subdirectories");
// Act — list the root's direct child subdirectories
var result = await InvokeToolAsync(listSubdirectories, new AIFunctionArguments());
// Assert — only direct children (reports, data); not the nested 2024
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().Select(e => e.GetString()).ToList();
Assert.Equal(2, entries.Count);
Assert.Contains("reports", entries);
Assert.Contains("data", entries);
}
[Fact]
public async Task ListSubdirectories_WithDirectory_ListsNestedChildrenAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("reports/q1.md", "Q1");
await store.WriteFileAsync("reports/2024/q2.md", "Q2");
await store.WriteFileAsync("reports/2025/q3.md", "Q3");
var tools = await CreateToolsAsync(store);
var listSubdirectories = GetTool(tools, "file_access_list_subdirectories");
// Act
var result = await InvokeToolAsync(listSubdirectories, new AIFunctionArguments
{
["directory"] = "reports",
});
// Assert — direct child subdirectories of reports
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().Select(e => e.GetString()).ToList();
Assert.Equal(2, entries.Count);
Assert.Contains("2024", entries);
Assert.Contains("2025", entries);
}
[Fact]
public async Task ListSubdirectories_NoSubdirectories_ReturnsEmptyAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("a.txt", "A");
await store.WriteFileAsync("b.txt", "B");
var tools = await CreateToolsAsync(store);
var listSubdirectories = GetTool(tools, "file_access_list_subdirectories");
// Act
var result = await InvokeToolAsync(listSubdirectories, new AIFunctionArguments());
// Assert
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
Assert.Empty(entries);
}
#endregion
#region SearchFiles Tests
@@ -459,7 +367,7 @@ public class FileAccessProviderTests
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Important research findings about AI");
var tools = await CreateToolsAsync(store);
var searchFiles = GetTool(tools, "file_access_search_files");
var searchFiles = GetTool(tools, "FileAccess_SearchFiles");
// Act
var result = await InvokeToolAsync(searchFiles, new AIFunctionArguments
@@ -484,7 +392,7 @@ public class FileAccessProviderTests
await store.WriteFileAsync("notes.md", "Important data");
await store.WriteFileAsync("data.txt", "Important data");
var tools = await CreateToolsAsync(store);
var searchFiles = GetTool(tools, "file_access_search_files");
var searchFiles = GetTool(tools, "FileAccess_SearchFiles");
// Act
var result = await InvokeToolAsync(searchFiles, new AIFunctionArguments
@@ -506,7 +414,7 @@ public class FileAccessProviderTests
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "No matching content here");
var tools = await CreateToolsAsync(store);
var searchFiles = GetTool(tools, "file_access_search_files");
var searchFiles = GetTool(tools, "FileAccess_SearchFiles");
// Act
var result = await InvokeToolAsync(searchFiles, new AIFunctionArguments
@@ -519,84 +427,6 @@ public class FileAccessProviderTests
Assert.Empty(entries);
}
[Fact]
public async Task SearchFiles_SearchesAllDescendantsRecursivelyAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("root.md", "Important data at root");
await store.WriteFileAsync("reports/q1.md", "Important data in reports");
await store.WriteFileAsync("reports/2024/q2.md", "Important data nested deeper");
var tools = await CreateToolsAsync(store);
var searchFiles = GetTool(tools, "file_access_search_files");
// Act — no glob, so all descendants are searched
var result = await InvokeToolAsync(searchFiles, new AIFunctionArguments
{
["regexPattern"] = "Important",
});
// Assert — matches at every depth, returned as store-root-relative paths
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
var names = entries.ConvertAll(e => e.GetProperty("fileName").GetString());
Assert.Equal(3, names.Count);
Assert.Contains("root.md", names);
Assert.Contains("reports/q1.md", names);
Assert.Contains("reports/2024/q2.md", names);
}
[Fact]
public async Task SearchFiles_GlobScopesToSubtreeAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("root.md", "Important data at root");
await store.WriteFileAsync("reports/q1.md", "Important data in reports");
await store.WriteFileAsync("reports/2024/q2.md", "Important data nested deeper");
var tools = await CreateToolsAsync(store);
var searchFiles = GetTool(tools, "file_access_search_files");
// Act — restrict to the reports subtree using a recursive glob
var result = await InvokeToolAsync(searchFiles, new AIFunctionArguments
{
["regexPattern"] = "Important",
["filePattern"] = "reports/**",
});
// Assert — only the files under reports/ match
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
var names = entries.ConvertAll(e => e.GetProperty("fileName").GetString());
Assert.Equal(2, names.Count);
Assert.Contains("reports/q1.md", names);
Assert.Contains("reports/2024/q2.md", names);
}
[Fact]
public async Task SearchFiles_RecursiveGlobMatchesNestedExtensionAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Important data");
await store.WriteFileAsync("data/raw.txt", "Important data");
await store.WriteFileAsync("reports/2024/q1.md", "Important data");
var tools = await CreateToolsAsync(store);
var searchFiles = GetTool(tools, "file_access_search_files");
// Act — match markdown files at any depth
var result = await InvokeToolAsync(searchFiles, new AIFunctionArguments
{
["regexPattern"] = "Important",
["filePattern"] = "**/*.md",
});
// Assert
var entries = Assert.IsType<JsonElement>(result).EnumerateArray().ToList();
var names = entries.ConvertAll(e => e.GetProperty("fileName").GetString());
Assert.Equal(2, names.Count);
Assert.Contains("notes.md", names);
Assert.Contains("reports/2024/q1.md", names);
}
#endregion
#region Path Traversal Protection
@@ -606,7 +436,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(async () =>
@@ -622,7 +452,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(async () =>
@@ -638,7 +468,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(async () =>
@@ -655,7 +485,7 @@ public class FileAccessProviderTests
// Arrange — "notes..md" is not a path traversal attempt.
var store = new InMemoryAgentFileStore();
var tools = await CreateToolsAsync(store);
var saveFile = GetTool(tools, "file_access_save_file");
var saveFile = GetTool(tools, "FileAccess_SaveFile");
// Act
await InvokeToolAsync(saveFile, new AIFunctionArguments
@@ -673,7 +503,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var readFile = GetTool(tools, "file_access_read_file");
var readFile = GetTool(tools, "FileAccess_ReadFile");
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(async () =>
@@ -688,7 +518,7 @@ public class FileAccessProviderTests
{
// Arrange
var tools = await CreateToolsAsync();
var deleteFile = GetTool(tools, "file_access_delete_file");
var deleteFile = GetTool(tools, "FileAccess_DeleteFile");
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(async () =>
@@ -2,7 +2,6 @@
using System;
using System.IO;
using System.Linq;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
@@ -334,107 +333,6 @@ public sealed class FileSystemAgentFileStoreTests : IDisposable
this._store.SearchFilesAsync("", "(a+)+$"));
}
[Fact]
public async Task SearchFilesAsync_Recursive_FindsDescendantsAsync()
{
// Arrange
await this._store.WriteFileAsync("notes.md", "Match here");
await this._store.WriteFileAsync("reports/q1.md", "Match here too");
await this._store.WriteFileAsync("reports/2024/q2.md", "Match here as well");
// Act
var results = await this._store.SearchFilesAsync("", "Match", filePattern: null, recursive: true);
// Assert
Assert.Equal(3, results.Count);
var names = string.Join(",", results.Select(r => r.FileName).OrderBy(n => n, StringComparer.Ordinal));
Assert.Equal("notes.md,reports/2024/q2.md,reports/q1.md", names);
}
[Fact]
public async Task SearchFilesAsync_Recursive_GlobScopesToSubtreeAsync()
{
// Arrange
await this._store.WriteFileAsync("notes.md", "Match here");
await this._store.WriteFileAsync("reports/q1.md", "Match here too");
await this._store.WriteFileAsync("reports/2024/q2.md", "Match here as well");
// Act
var results = await this._store.SearchFilesAsync("", "Match", filePattern: "reports/**", recursive: true);
// Assert
Assert.Equal(2, results.Count);
var names = string.Join(",", results.Select(r => r.FileName).OrderBy(n => n, StringComparer.Ordinal));
Assert.Equal("reports/2024/q2.md,reports/q1.md", names);
}
[Fact]
public async Task SearchFilesAsync_Recursive_GlobMatchesNestedExtensionAsync()
{
// Arrange
await this._store.WriteFileAsync("notes.md", "Match here");
await this._store.WriteFileAsync("reports/q1.txt", "Match here too");
await this._store.WriteFileAsync("reports/2024/q2.md", "Match here as well");
// Act
var results = await this._store.SearchFilesAsync("", "Match", filePattern: "**/*.md", recursive: true);
// Assert
Assert.Equal(2, results.Count);
var names = string.Join(",", results.Select(r => r.FileName).OrderBy(n => n, StringComparer.Ordinal));
Assert.Equal("notes.md,reports/2024/q2.md", names);
}
[Fact]
public async Task ListDirectoriesAsync_ReturnsDirectChildSubdirectoriesAsync()
{
// Arrange
await this._store.WriteFileAsync("root.md", "x");
await this._store.WriteFileAsync("reports/q1.md", "x");
await this._store.WriteFileAsync("reports/2024/q2.md", "x");
await this._store.WriteFileAsync("images/logo.txt", "x");
// Act
var directories = await this._store.ListDirectoriesAsync("");
// Assert
var sorted = string.Join(",", directories.OrderBy(d => d, StringComparer.Ordinal));
Assert.Equal("images,reports", sorted);
}
[Fact]
public async Task ListDirectoriesAsync_NestedDirectory_ReturnsChildrenAsync()
{
// Arrange
await this._store.WriteFileAsync("reports/q1.md", "x");
await this._store.WriteFileAsync("reports/2024/q2.md", "x");
await this._store.WriteFileAsync("reports/2025/q3.md", "x");
// Act
var directories = await this._store.ListDirectoriesAsync("reports");
// Assert
var sorted = string.Join(",", directories.OrderBy(d => d, StringComparer.Ordinal));
Assert.Equal("2024,2025", sorted);
}
[Fact]
public async Task ListDirectoriesAsync_NonExistentDirectory_ReturnsEmptyAsync()
{
// Act
var directories = await this._store.ListDirectoriesAsync("no-dir");
// Assert
Assert.Empty(directories);
}
[Fact]
public async Task ListDirectoriesAsync_DotDotSegment_ThrowsAsync()
{
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(() => this._store.ListDirectoriesAsync("../other"));
}
#endregion
#region Symlink Escape Rejection
@@ -904,79 +802,6 @@ public sealed class FileSystemAgentFileStoreTests : IDisposable
File.Delete(outsideFile);
}
}
[Fact]
public async Task SearchFilesAsync_Recursive_SkipsSymlinkedSubdirectoryAsync()
{
// Arrange — a symlinked directory under root should be skipped by recursive search.
string outsideDir = Path.Combine(Path.GetTempPath(), "symlink_recursive_target_" + Guid.NewGuid().ToString("N"));
Directory.CreateDirectory(outsideDir);
File.WriteAllText(Path.Combine(outsideDir, "leak.txt"), "RECURSIVE_SECRET_CONTENT");
string linkDir = Path.Combine(this._rootDir, "linked-sub");
try
{
if (!TryCreateDirectorySymbolicLink(linkDir, outsideDir))
{
return;
}
await this._store.WriteFileAsync("normal/visible.txt", "RECURSIVE_VISIBLE_CONTENT");
// Act — recursive search should not descend into the symlinked directory.
var results = await this._store.SearchFilesAsync("", "RECURSIVE", filePattern: null, recursive: true);
// Assert — only the non-symlinked file is found.
Assert.Single(results);
Assert.Equal("normal/visible.txt", results[0].FileName);
}
finally
{
if (Directory.Exists(linkDir))
{
Directory.Delete(linkDir);
}
Directory.Delete(outsideDir, recursive: true);
}
}
[Fact]
public async Task ListDirectoriesAsync_ExcludesSymlinkedDirectoryAsync()
{
// Arrange — a symlinked directory under root should not be listed.
string outsideDir = Path.Combine(Path.GetTempPath(), "symlink_listdir_target_" + Guid.NewGuid().ToString("N"));
Directory.CreateDirectory(outsideDir);
string linkDir = Path.Combine(this._rootDir, "linked-listing");
try
{
if (!TryCreateDirectorySymbolicLink(linkDir, outsideDir))
{
return;
}
await this._store.WriteFileAsync("real-dir/file.txt", "x");
// Act
var directories = await this._store.ListDirectoriesAsync("");
// Assert — the symlinked directory is excluded, the real one is present.
Assert.DoesNotContain("linked-listing", directories);
Assert.Contains("real-dir", directories);
}
finally
{
if (Directory.Exists(linkDir))
{
Directory.Delete(linkDir);
}
Directory.Delete(outsideDir, recursive: true);
}
}
#endif
#endregion
@@ -1,7 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Linq;
using System.Threading.Tasks;
namespace Microsoft.Agents.AI.UnitTests.Harness.FileMemory;
@@ -521,117 +520,4 @@ public class InMemoryAgentFileStoreTests
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(() => store.ListFilesAsync("../other"));
}
[Fact]
public async Task ListDirectories_PathTraversal_ThrowsAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
// Act & Assert
await Assert.ThrowsAsync<ArgumentException>(() => store.ListDirectoriesAsync("../other"));
}
[Fact]
public async Task SearchFiles_Recursive_FindsDescendantsAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Match here");
await store.WriteFileAsync("reports/q1.md", "Match here too");
await store.WriteFileAsync("reports/2024/q2.md", "Match here as well");
// Act
var results = await store.SearchFilesAsync("", "Match", filePattern: null, recursive: true);
// Assert
Assert.Equal(3, results.Count);
var names = string.Join(",", results.Select(r => r.FileName).OrderBy(n => n, StringComparer.Ordinal));
Assert.Equal("notes.md,reports/2024/q2.md,reports/q1.md", names);
}
[Fact]
public async Task SearchFiles_Recursive_GlobScopesToSubtreeAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Match here");
await store.WriteFileAsync("reports/q1.md", "Match here too");
await store.WriteFileAsync("reports/2024/q2.md", "Match here as well");
// Act
var results = await store.SearchFilesAsync("", "Match", filePattern: "reports/**", recursive: true);
// Assert
Assert.Equal(2, results.Count);
var names = string.Join(",", results.Select(r => r.FileName).OrderBy(n => n, StringComparer.Ordinal));
Assert.Equal("reports/2024/q2.md,reports/q1.md", names);
}
[Fact]
public async Task SearchFiles_Recursive_GlobMatchesNestedExtensionAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("notes.md", "Match here");
await store.WriteFileAsync("reports/q1.txt", "Match here too");
await store.WriteFileAsync("reports/2024/q2.md", "Match here as well");
// Act
var results = await store.SearchFilesAsync("", "Match", filePattern: "**/*.md", recursive: true);
// Assert
Assert.Equal(2, results.Count);
var names = string.Join(",", results.Select(r => r.FileName).OrderBy(n => n, StringComparer.Ordinal));
Assert.Equal("notes.md,reports/2024/q2.md", names);
}
[Fact]
public async Task ListDirectories_ReturnsDirectChildSubdirectoriesAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("root.md", "x");
await store.WriteFileAsync("reports/q1.md", "x");
await store.WriteFileAsync("reports/2024/q2.md", "x");
await store.WriteFileAsync("images/logo.png", "x");
// Act
var directories = await store.ListDirectoriesAsync("");
// Assert
var sorted = string.Join(",", directories.OrderBy(d => d, StringComparer.Ordinal));
Assert.Equal("images,reports", sorted);
}
[Fact]
public async Task ListDirectories_NestedDirectory_ReturnsChildrenAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("reports/q1.md", "x");
await store.WriteFileAsync("reports/2024/q2.md", "x");
await store.WriteFileAsync("reports/2025/q3.md", "x");
// Act
var directories = await store.ListDirectoriesAsync("reports");
// Assert
var sorted = string.Join(",", directories.OrderBy(d => d, StringComparer.Ordinal));
Assert.Equal("2024,2025", sorted);
}
[Fact]
public async Task ListDirectories_NoSubdirectories_ReturnsEmptyAsync()
{
// Arrange
var store = new InMemoryAgentFileStore();
await store.WriteFileAsync("root.md", "x");
// Act
var directories = await store.ListDirectoriesAsync("");
// Assert
Assert.Empty(directories);
}
}
-1
View File
@@ -1 +0,0 @@
../../../.github/skills/pull-requests
-1
View File
@@ -14,7 +14,6 @@ Instructions for AI coding agents working in the Python codebase.
- `python-feature-lifecycle` — package vs feature lifecycle stages, decorators, enums, and promotion guidance
- `python-package-management` — monorepo structure, lazy loading, versioning, new packages
- `python-samples` — sample file structure, PEP 723, documentation guidelines
- `pull-requests` — writing PR descriptions (template) and handling/resolving PR review comments
## Maintaining Documentation
-2
View File
@@ -10,12 +10,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard.
- **`AGUIHttpService`** - HTTP service for AG-UI endpoints
- **`AGUIEventConverter`** - Converts between Agent Framework and AG-UI events
- **`add_agent_framework_fastapi_endpoint()`** - Add AG-UI endpoint to FastAPI app (`SupportsAgentRun` or `Workflow`)
- **`InMemoryAGUIThreadSnapshotStore`** - Memory-only latest AG-UI Thread Snapshot store for local development, demos, and tests
## Types
- **`AGUIRequest`** / **`AGUIChatOptions`** - Request types
- **`AGUIThreadSnapshot`** / **`AGUIThreadSnapshotStore`** - Replayable thread snapshot model and scoped async store protocol
- **`availableInterrupts` / `resume`** - Optional interrupt configuration and continuation payloads
- **`AgentState`** / **`RunMetadata`** - State management types
- **`PredictStateConfig`** - Configuration for state prediction
-65
View File
@@ -198,71 +198,6 @@ The `dependencies` parameter accepts any FastAPI dependency, enabling integratio
For a complete authentication example, see [getting_started/server.py](getting_started/server.py).
## AG-UI Thread Snapshots
AG-UI Thread Snapshot persistence is opt-in and disabled by default. Existing endpoints keep their current behavior
unless you provide a `snapshot_store`.
Thread snapshots let an AG-UI frontend recover replayable UI state after a refresh. When snapshot persistence is
enabled, the endpoint stores the latest replayable snapshot for an AG-UI Thread within an application-defined
Snapshot Scope. A Hydrate Request is an AG-UI request with a known `threadId`, `messages: []`, and no `resume`
payload. Hydration replays the stored Shared State, message snapshot, and interruption metadata when available,
then finishes without invoking the wrapped agent or workflow.
Use the built-in in-memory store for local development, demos, and tests:
```python
from fastapi import FastAPI
from agent_framework.ag_ui import InMemoryAGUIThreadSnapshotStore, add_agent_framework_fastapi_endpoint
app = FastAPI()
agent = ...
snapshot_store = InMemoryAGUIThreadSnapshotStore(max_snapshots=500)
def resolve_snapshot_scope(request):
# Local demo scope. Production apps should derive the scope from authenticated user or tenant context.
del request
return "local-demo"
add_agent_framework_fastapi_endpoint(
app,
agent,
"/",
snapshot_store=snapshot_store,
snapshot_scope_resolver=resolve_snapshot_scope,
)
```
A frontend can then hydrate the latest stored snapshot for the scoped thread:
```json
{
"threadId": "thread-1",
"messages": []
}
```
Endpoint configuration requires `snapshot_scope_resolver` whenever a snapshot store is configured, including when
the store is already set on a pre-wrapped `AgentFrameworkAgent` or `AgentFrameworkWorkflow`. The resolver returns
the application-defined Snapshot Scope used with the AG-UI Thread id as the storage key.
AG-UI Thread ids identify AG-UI Threads; they do not authorize snapshot access. Do not treat a thread id as a bearer
credential or tenant boundary. Production applications must authenticate and authorize every AG-UI endpoint request
and choose a Snapshot Scope that represents the app's real access boundary, such as an authenticated user, tenant,
or workspace. Do not rely on untrusted client-provided fields by themselves to choose that boundary.
Stored snapshots are untrusted application data with confidentiality impact. They may contain sensitive user text,
model output, tool results, function arguments, UI payloads, Shared State, and interruption data. The built-in
`InMemoryAGUIThreadSnapshotStore` is in-memory only, process-local, bounded, latest-only, and not durable production
storage. It is cleared on process restart and is not shared across workers.
No file-backed AG-UI snapshot store is provided by the package. Applications that need durable persistence should
provide an app-owned implementation of the `AGUIThreadSnapshotStore` protocol and own storage hardening, including
encryption, access control, retention, audit, data residency, and deletion behavior.
## Architecture
The package uses a clean, orchestrator-based architecture:
@@ -9,15 +9,6 @@ from ._client import AGUIChatClient
from ._endpoint import add_agent_framework_fastapi_endpoint
from ._event_converters import AGUIEventConverter
from ._http_service import AGUIHttpService
from ._snapshots import (
DEFAULT_MAX_THREAD_SNAPSHOTS,
AGUIThreadID,
AGUIThreadSnapshot,
AGUIThreadSnapshotStore,
InMemoryAGUIThreadSnapshotStore,
SnapshotScope,
SnapshotScopeResolver,
)
from ._state import state_update
from ._types import AgentState, AGUIChatOptions, AGUIRequest, PredictStateConfig, RunMetadata
from ._workflow import AgentFrameworkWorkflow, WorkflowFactory
@@ -40,16 +31,9 @@ __all__ = [
"AGUIEventConverter",
"AGUIHttpService",
"AGUIRequest",
"AGUIThreadID",
"AGUIThreadSnapshot",
"AGUIThreadSnapshotStore",
"AgentState",
"InMemoryAGUIThreadSnapshotStore",
"PredictStateConfig",
"RunMetadata",
"SnapshotScope",
"SnapshotScopeResolver",
"DEFAULT_MAX_THREAD_SNAPSHOTS",
"DEFAULT_TAGS",
"state_update",
"__version__",
@@ -10,7 +10,6 @@ from ag_ui.core import BaseEvent
from agent_framework import SupportsAgentRun
from ._agent_run import PendingApprovalEntry, run_agent_stream
from ._snapshots import AGUIThreadSnapshotStore
class AgentConfig:
@@ -22,7 +21,6 @@ class AgentConfig:
predict_state_config: dict[str, dict[str, str]] | None = None,
use_service_session: bool = False,
require_confirmation: bool = True,
snapshot_store: AGUIThreadSnapshotStore | None = None,
):
"""Initialize agent configuration.
@@ -31,14 +29,11 @@ class AgentConfig:
predict_state_config: Configuration for predictive state updates
use_service_session: Whether the agent session is service-managed
require_confirmation: Whether predictive updates require user confirmation before applying
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
endpoint setup also provides an explicit Snapshot Scope resolver.
"""
self.state_schema = self._normalize_state_schema(state_schema)
self.predict_state_config = predict_state_config or {}
self.use_service_session = use_service_session
self.require_confirmation = require_confirmation
self.snapshot_store = snapshot_store
@staticmethod
def _normalize_state_schema(state_schema: Any | None) -> dict[str, Any]:
@@ -84,7 +79,6 @@ class AgentFrameworkAgent:
predict_state_config: dict[str, dict[str, str]] | None = None,
require_confirmation: bool = True,
use_service_session: bool = False,
snapshot_store: AGUIThreadSnapshotStore | None = None,
):
"""Initialize the AG-UI compatible agent wrapper.
@@ -96,8 +90,6 @@ class AgentFrameworkAgent:
predict_state_config: Configuration for predictive state updates
require_confirmation: Whether predictive updates require user confirmation before applying
use_service_session: Whether the agent session is service-managed
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
endpoint setup also provides an explicit Snapshot Scope resolver.
"""
self.agent = agent
self.name = name or getattr(agent, "name", "agent")
@@ -108,7 +100,6 @@ class AgentFrameworkAgent:
predict_state_config=predict_state_config,
use_service_session=use_service_session,
require_confirmation=require_confirmation,
snapshot_store=snapshot_store,
)
# Server-side registry of pending approval requests.
@@ -119,11 +110,6 @@ class AgentFrameworkAgent:
self._pending_approvals: OrderedDict[str, PendingApprovalEntry] = OrderedDict()
self._pending_approvals_max_size: int = 10_000
@property
def snapshot_store(self) -> AGUIThreadSnapshotStore | None:
"""Configured AG-UI Thread Snapshot store, if any."""
return self.config.snapshot_store
async def run(
self,
input_data: dict[str, Any],
@@ -4,7 +4,6 @@
from __future__ import annotations # noqa: I001
import copy
import json
import logging
import uuid
@@ -53,11 +52,9 @@ from ._run_common import (
_extract_tool_result_display, # type: ignore
_has_only_tool_calls, # type: ignore
_normalize_resume_interrupts, # type: ignore
_reconstruct_messages_from_thread_snapshot, # type: ignore
_resolve_ui_payload, # type: ignore
_stringify_tool_result, # type: ignore
)
from ._snapshots import AGUIThreadSnapshot, _DEFAULT_STATE_INPUT_KEY, _SNAPSHOT_SCOPE_INPUT_KEY
from ._utils import (
canonical_function_arguments,
convert_agui_tools_to_agent_framework,
@@ -751,85 +748,6 @@ def _build_messages_snapshot(
return MessagesSnapshotEvent(messages=all_messages) # type: ignore[arg-type]
def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]:
"""Convert AG-UI message event models back to plain snapshot dictionaries."""
safe_messages = make_json_safe(messages)
if not isinstance(safe_messages, list):
return []
return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)]
def _text_events_to_snapshot_messages(events: list[BaseEvent]) -> list[dict[str, Any]]:
"""Convert streamed text-message events into snapshot message dictionaries."""
messages: list[dict[str, Any]] = []
messages_by_id: dict[str, dict[str, Any]] = {}
for event in events:
if isinstance(event, TextMessageStartEvent):
message: dict[str, Any] = {"id": event.message_id, "role": event.role, "content": ""}
messages.append(message)
messages_by_id[event.message_id] = message
elif isinstance(event, TextMessageContentEvent):
open_message = messages_by_id.get(event.message_id)
if open_message is not None:
open_message["content"] = f"{open_message['content']}{event.delta}"
return [message for message in messages if message.get("content")]
async def _hydrate_thread_snapshot(
*,
config: AgentConfig,
scope: str,
thread_id: str,
run_id: str,
) -> AsyncGenerator[BaseEvent]:
"""Replay the latest stored AG-UI Thread Snapshot without invoking the agent."""
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
if config.snapshot_store is None:
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
return
snapshot = await config.snapshot_store.get(scope=scope, thread_id=thread_id)
if snapshot is None:
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
return
if snapshot.state is not None:
yield StateSnapshotEvent(snapshot=snapshot.state)
if snapshot.messages:
yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type]
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=snapshot.interrupt)
async def _save_thread_snapshot(
*,
config: AgentConfig,
scope: str | None,
thread_id: str,
messages: list[dict[str, Any]],
state: dict[str, Any] | None,
interrupt: list[dict[str, Any]] | None,
) -> None:
"""Save the latest replayable AG-UI Thread Snapshot when persistence is configured."""
if config.snapshot_store is None or scope is None:
return
try:
await config.snapshot_store.save(
scope=scope,
thread_id=thread_id,
snapshot=AGUIThreadSnapshot(messages=messages, state=state, interrupt=interrupt),
)
except Exception:
# The run itself already streamed successfully; a transient store failure
# must not surface as RUN_ERROR for a completed run. The previous snapshot
# stays available for hydration.
logger.exception(
"Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.",
scope,
thread_id,
)
async def run_agent_stream(
input_data: dict[str, Any],
agent: SupportsAgentRun,
@@ -856,53 +774,15 @@ async def run_agent_stream(
# Parse IDs
thread_id = input_data.get("thread_id") or input_data.get("threadId") or str(uuid.uuid4())
run_id = input_data.get("run_id") or input_data.get("runId") or str(uuid.uuid4())
snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY))
# Initialize flow state with schema defaults
flow = FlowState()
if input_data.get("state"):
flow.current_state = dict(input_data["state"])
state_schema = cast(dict[str, Any], getattr(config, "state_schema", {}) or {})
predict_state_config = cast(dict[str, dict[str, str]], getattr(config, "predict_state_config", {}) or {})
# Normalize messages
available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts")
raw_messages: list[dict[str, Any]] = input_data.get("messages", []) or []
resume_payload = _extract_resume_payload(input_data)
if config.snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None:
async for event in _hydrate_thread_snapshot(
config=config,
scope=snapshot_scope,
thread_id=thread_id,
run_id=run_id,
):
yield event
return
stored_snapshot: AGUIThreadSnapshot | None = None
if config.snapshot_store is not None and snapshot_scope is not None:
stored_snapshot = await config.snapshot_store.get(scope=snapshot_scope, thread_id=thread_id)
if stored_snapshot is not None and resume_payload is None:
raw_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages,
incoming_messages=raw_messages,
stored_interrupt=stored_snapshot.interrupt,
)
# Initialize flow state with stored state plus request-provided overrides.
flow = FlowState()
request_state = input_data.get("state")
if stored_snapshot is not None and stored_snapshot.state is not None:
flow.current_state = dict(stored_snapshot.state)
if isinstance(request_state, dict):
flow.current_state.update(request_state)
elif isinstance(request_state, dict):
flow.current_state = dict(request_state)
# Apply endpoint-deferred defaults only for keys missing from both the stored
# snapshot state and the request state, so defaults never reset persisted state.
deferred_default_state = cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY))
if deferred_default_state:
for key, value in deferred_default_state.items():
if key not in flow.current_state:
flow.current_state[key] = copy.deepcopy(value)
# Apply schema defaults for missing state keys
if state_schema:
for key, schema in state_schema.items():
@@ -921,7 +801,10 @@ async def run_agent_stream(
current_state=flow.current_state,
)
resume_messages = _resume_to_tool_messages(resume_payload)
# Normalize messages
available_interrupts = input_data.get("available_interrupts") or input_data.get("availableInterrupts")
raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or []))
resume_messages = _resume_to_tool_messages(_extract_resume_payload(input_data))
if available_interrupts:
logger.debug("Received available interrupts metadata: %s", available_interrupts)
if resume_messages:
@@ -1009,24 +892,8 @@ async def run_agent_stream(
# Emit approved state snapshot before confirmation message
if approved_state_snapshot_emitted:
yield StateSnapshotEvent(snapshot=flow.current_state)
confirmation_events = _handle_step_based_approval(messages)
for event in confirmation_events:
for event in _handle_step_based_approval(messages):
yield event
# Persist the completed confirmation turn with interrupt=None so hydration
# does not replay the stale pending interrupt after the user responded.
persisted_messages = snapshot_messages + _text_events_to_snapshot_messages(confirmation_events)
if resume_payload is not None and stored_snapshot is not None:
# Resume requests carry only the synthesized interrupt response, so prepend
# the stored thread history to avoid persisting a truncated thread.
persisted_messages = [copy.deepcopy(message) for message in stored_snapshot.messages] + persisted_messages
await _save_thread_snapshot(
config=config,
scope=snapshot_scope,
thread_id=thread_id,
messages=persisted_messages,
state=cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None,
interrupt=None,
)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
return
@@ -1038,9 +905,6 @@ async def run_agent_stream(
# Stream from agent - emit RunStarted after first update to get service IDs
run_started_emitted = False
all_updates: list[Any] = [] # Collect for structured output processing
latest_state_snapshot: dict[str, Any] | None = (
cast(dict[str, Any], make_json_safe(flow.current_state)) if flow.current_state else None
)
response_stream = agent.run(messages, stream=True, **run_kwargs)
stream = await _normalize_response_stream(response_stream)
async for update in stream:
@@ -1070,7 +934,6 @@ async def run_agent_stream(
yield CustomEvent(name="PredictState", value=predict_state_value)
# Emit initial state snapshot only if we have both state_schema and state
if state_schema and flow.current_state:
latest_state_snapshot = cast(dict[str, Any], make_json_safe(flow.current_state))
yield StateSnapshotEvent(snapshot=flow.current_state)
run_started_emitted = True
@@ -1112,8 +975,6 @@ async def run_agent_stream(
skip_text,
config.require_confirmation,
):
if isinstance(event, StateSnapshotEvent):
latest_state_snapshot = cast(dict[str, Any], make_json_safe(event.snapshot))
yield event
# Stop if waiting for approval
@@ -1158,7 +1019,6 @@ async def run_agent_stream(
if state_updates:
flow.current_state.update(state_updates)
latest_state_snapshot = cast(dict[str, Any], make_json_safe(flow.current_state))
yield StateSnapshotEvent(snapshot=flow.current_state)
logger.info(f"Emitted StateSnapshotEvent with updates: {list(state_updates.keys())}")
@@ -1196,7 +1056,6 @@ async def run_agent_stream(
if result:
state_key, state_value = result
flow.current_state[state_key] = state_value
latest_state_snapshot = cast(dict[str, Any], make_json_safe(flow.current_state))
yield StateSnapshotEvent(snapshot=flow.current_state)
except json.JSONDecodeError:
# Ignore malformed JSON in tool arguments for predictive state;
@@ -1277,12 +1136,7 @@ async def run_agent_stream(
should_emit_snapshot = (
flow.pending_tool_calls or flow.tool_results or flow.accumulated_text or flow.reasoning_messages
)
latest_messages_snapshot = snapshot_messages
if should_emit_snapshot:
# Always fold this turn's output into the persisted snapshot, even when the
# outbound MESSAGES_SNAPSHOT event is suppressed for predictive tools.
snapshot_event = _build_messages_snapshot(flow, snapshot_messages)
latest_messages_snapshot = _event_messages_to_snapshot_dicts(list(snapshot_event.messages))
# Check if we should suppress for predictive tool
last_tool_name = None
if flow.tool_results:
@@ -1292,21 +1146,8 @@ async def run_agent_stream(
if not _should_suppress_intermediate_snapshot(
last_tool_name, predict_state_config, config.require_confirmation
):
yield snapshot_event
yield _build_messages_snapshot(flow, snapshot_messages)
# Always emit RunFinished - confirm_changes tool call is complete (Start -> Args -> End)
# The UI will show confirmation dialog and send a new request when user responds
persisted_messages = latest_messages_snapshot
if resume_payload is not None and stored_snapshot is not None:
# Resume requests carry only the synthesized interrupt response, so prepend
# the stored thread history to avoid persisting a truncated thread.
persisted_messages = [copy.deepcopy(message) for message in stored_snapshot.messages] + persisted_messages
await _save_thread_snapshot(
config=config,
scope=snapshot_scope,
thread_id=thread_id,
messages=persisted_messages,
state=latest_state_snapshot,
interrupt=flow.interrupts or None,
)
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=flow.interrupts)
@@ -7,7 +7,6 @@ from __future__ import annotations
import copy
import logging
from collections.abc import AsyncGenerator, Sequence
from inspect import isawaitable
from typing import Any
from ag_ui.core import RunErrorEvent
@@ -18,58 +17,12 @@ from fastapi.params import Depends
from fastapi.responses import StreamingResponse
from ._agent import AgentFrameworkAgent
from ._snapshots import (
_DEFAULT_STATE_INPUT_KEY,
_SNAPSHOT_SCOPE_INPUT_KEY,
AGUIThreadSnapshotStore,
SnapshotScopeResolver,
)
from ._types import AGUIRequest
from ._workflow import AgentFrameworkWorkflow
logger = logging.getLogger(__name__)
def _get_snapshot_store(
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
) -> AGUIThreadSnapshotStore | None:
if isinstance(protocol_runner, AgentFrameworkAgent):
return protocol_runner.config.snapshot_store
return protocol_runner.snapshot_store
def _set_snapshot_store(
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
snapshot_store: AGUIThreadSnapshotStore,
) -> None:
if isinstance(protocol_runner, AgentFrameworkAgent):
protocol_runner.config.snapshot_store = snapshot_store
return
protocol_runner.snapshot_store = snapshot_store
def _configure_snapshot_persistence(
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow,
*,
snapshot_store: AGUIThreadSnapshotStore | None,
snapshot_scope_resolver: SnapshotScopeResolver | None,
) -> None:
existing_snapshot_store = _get_snapshot_store(protocol_runner)
if snapshot_store is not None:
if existing_snapshot_store is not None and existing_snapshot_store is not snapshot_store:
raise ValueError("snapshot_store is already configured on the AG-UI runner.")
if existing_snapshot_store is None:
_set_snapshot_store(protocol_runner, snapshot_store)
existing_snapshot_store = snapshot_store
if existing_snapshot_store is not None and snapshot_scope_resolver is None:
raise ValueError(
"snapshot_scope_resolver is required when snapshot_store is configured. "
"AG-UI Thread ids identify threads but do not authorize snapshot access; "
"provide a resolver that returns an explicit Snapshot Scope."
)
def add_agent_framework_fastapi_endpoint(
app: FastAPI,
agent: SupportsAgentRun | AgentFrameworkAgent | Workflow | AgentFrameworkWorkflow,
@@ -80,8 +33,6 @@ def add_agent_framework_fastapi_endpoint(
default_state: dict[str, Any] | None = None,
tags: list[str] | None = None,
dependencies: Sequence[Depends] | None = None,
snapshot_store: AGUIThreadSnapshotStore | None = None,
snapshot_scope_resolver: SnapshotScopeResolver | None = None,
) -> None:
"""Add an AG-UI endpoint to a FastAPI app.
@@ -99,10 +50,6 @@ def add_agent_framework_fastapi_endpoint(
These dependencies run before the endpoint handler. Use this to add
authentication checks, rate limiting, or other middleware-like behavior.
Example: `dependencies=[Depends(verify_api_key)]`
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence is opt-in and requires an
explicit Snapshot Scope resolver.
snapshot_scope_resolver: Optional resolver for the application-defined Snapshot Scope. Required whenever
a snapshot store is configured because an AG-UI Thread id is not an authorization boundary.
"""
protocol_runner: AgentFrameworkAgent | AgentFrameworkWorkflow
if isinstance(agent, AgentFrameworkWorkflow):
@@ -116,17 +63,10 @@ def add_agent_framework_fastapi_endpoint(
agent=agent,
state_schema=state_schema,
predict_state_config=predict_state_config,
snapshot_store=snapshot_store,
)
else:
raise TypeError("agent must be SupportsAgentRun, Workflow, AgentFrameworkAgent, or AgentFrameworkWorkflow.")
_configure_snapshot_persistence(
protocol_runner,
snapshot_store=snapshot_store,
snapshot_scope_resolver=snapshot_scope_resolver,
)
@app.post(path, tags=tags or ["AG-UI"], dependencies=dependencies, response_model=None) # type: ignore[arg-type]
async def agent_endpoint(request_body: AGUIRequest) -> StreamingResponse:
"""Handle AG-UI agent requests.
@@ -136,23 +76,11 @@ def add_agent_framework_fastapi_endpoint(
"""
try:
input_data = request_body.model_dump(exclude_none=True)
snapshot_persistence_active = False
if snapshot_scope_resolver is not None and _get_snapshot_store(protocol_runner) is not None:
snapshot_scope = snapshot_scope_resolver(request_body)
if isawaitable(snapshot_scope):
snapshot_scope = await snapshot_scope
input_data[_SNAPSHOT_SCOPE_INPUT_KEY] = snapshot_scope
snapshot_persistence_active = True
if default_state:
if snapshot_persistence_active:
# Defer default application to the runner so defaults only fill keys
# missing from both the stored snapshot state and the request state.
input_data[_DEFAULT_STATE_INPUT_KEY] = copy.deepcopy(default_state)
else:
state = input_data.setdefault("state", {})
for key, value in default_state.items():
if key not in state:
state[key] = copy.deepcopy(value)
state = input_data.setdefault("state", {})
for key, value in default_state.items():
if key not in state:
state[key] = copy.deepcopy(value)
logger.debug(
f"[{path}] Received request - Run ID: {input_data.get('run_id', 'no-run-id')}, "
f"Thread ID: {input_data.get('thread_id', 'no-thread-id')}, "
@@ -4,7 +4,6 @@
from __future__ import annotations
import copy
import json
import logging
from collections.abc import Mapping
@@ -34,7 +33,7 @@ from agent_framework import Content
from ._orchestration._predictive_state import PredictiveStateHandler
from ._state import TOOL_RESULT_DISPLAY_KEY, TOOL_RESULT_STATE_KEY
from ._utils import generate_event_id, make_json_safe, normalize_agui_role
from ._utils import generate_event_id, make_json_safe
logger = logging.getLogger(__name__)
@@ -734,117 +733,3 @@ def _emit_content(
return _emit_text_reasoning(content, flow)
logger.debug("Skipping unsupported content type in AG-UI emitter: %s", content_type)
return events
def _canonical_snapshot_message(message: dict[str, Any]) -> dict[str, Any]:
"""Normalize an AG-UI message for identity comparison without generated ids."""
from ._message_adapters import agui_messages_to_snapshot_format
normalized_message = agui_messages_to_snapshot_format([copy.deepcopy(message)])[0]
normalized_message.pop("id", None)
return cast(dict[str, Any], make_json_safe(normalized_message))
def _snapshot_messages_match(stored_message: dict[str, Any], incoming_message: dict[str, Any]) -> bool:
"""Return whether an incoming message already represents the stored snapshot message."""
stored_id = stored_message.get("id")
incoming_id = incoming_message.get("id")
if stored_id and incoming_id:
return str(stored_id) == str(incoming_id)
return _canonical_snapshot_message(stored_message) == _canonical_snapshot_message(incoming_message)
def _latest_user_message_index(messages: list[dict[str, Any]]) -> int | None:
"""Find the newest incoming user message index."""
for index in range(len(messages) - 1, -1, -1):
if normalize_agui_role(messages[index].get("role", "user")) == "user":
return index
return None
def _known_tool_call_ids(
stored_messages: list[dict[str, Any]],
stored_interrupt: list[dict[str, Any]] | None,
) -> set[str]:
"""Collect tool call ids the backend previously issued for this thread."""
known_ids: set[str] = set()
for message in stored_messages:
tool_calls = message.get("tool_calls") or message.get("toolCalls") or []
if not isinstance(tool_calls, list):
continue
for tool_call in cast(list[Any], tool_calls):
if isinstance(tool_call, dict):
tool_call_id = cast(dict[str, Any], tool_call).get("id")
if tool_call_id:
known_ids.add(str(tool_call_id))
for interrupt in stored_interrupt or []:
interrupt_id = interrupt.get("id")
if interrupt_id:
known_ids.add(str(interrupt_id))
return known_ids
def _filter_untrusted_suffix(
incoming_suffix: list[dict[str, Any]],
*,
stored_messages: list[dict[str, Any]],
stored_interrupt: list[dict[str, Any]] | None,
) -> list[dict[str, Any]]:
"""Drop client-forged non-user messages before promoting them to stored history.
Only the user's own turns and tool results answering backend-issued tool calls
(including pending interrupts) may extend the authoritative thread history.
"""
known_ids: set[str] | None = None
filtered: list[dict[str, Any]] = []
for message in incoming_suffix:
raw_role = str(message.get("role", "")).lower()
if raw_role == "user":
filtered.append(message)
continue
if raw_role == "tool":
tool_call_id = message.get("toolCallId") or message.get("tool_call_id") or message.get("actionExecutionId")
if known_ids is None:
known_ids = _known_tool_call_ids(stored_messages, stored_interrupt)
if tool_call_id and str(tool_call_id) in known_ids:
filtered.append(message)
continue
logger.warning(
"Dropping client-supplied %r message from the incoming thread suffix; "
"only user turns and tool results for backend-issued tool calls extend stored history.",
raw_role or "unknown",
)
return filtered
def _reconstruct_messages_from_thread_snapshot(
*,
stored_messages: list[dict[str, Any]],
incoming_messages: list[dict[str, Any]],
stored_interrupt: list[dict[str, Any]] | None = None,
) -> list[dict[str, Any]]:
"""Combine backend-owned prior history with the request-owned new user turn."""
if not stored_messages or not incoming_messages:
return incoming_messages
incoming_suffix: list[dict[str, Any]]
if len(incoming_messages) >= len(stored_messages) and all(
_snapshot_messages_match(stored_message, incoming_message)
for stored_message, incoming_message in zip(stored_messages, incoming_messages)
):
incoming_suffix = incoming_messages[len(stored_messages) :]
else:
latest_user_index = _latest_user_message_index(incoming_messages)
if latest_user_index is None:
return incoming_messages
incoming_suffix = incoming_messages[latest_user_index:]
incoming_suffix = _filter_untrusted_suffix(
incoming_suffix,
stored_messages=stored_messages,
stored_interrupt=stored_interrupt,
)
return [copy.deepcopy(message) for message in stored_messages] + [
copy.deepcopy(message) for message in incoming_suffix
]
@@ -1,202 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""AG-UI Thread Snapshot storage primitives."""
from __future__ import annotations
import copy
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Protocol, TypeAlias, runtime_checkable
if TYPE_CHECKING:
from ._types import AGUIRequest
SnapshotScope: TypeAlias = str
"""Application-defined scope for authorizing access to AG-UI Thread Snapshots."""
AGUIThreadID: TypeAlias = str
"""AG-UI Thread identifier within a Snapshot Scope."""
SnapshotScopeResolver: TypeAlias = Callable[["AGUIRequest"], str | Awaitable[str]]
"""Callable that resolves the Snapshot Scope for an AG-UI endpoint request."""
_SnapshotKey: TypeAlias = tuple[SnapshotScope, AGUIThreadID]
DEFAULT_MAX_THREAD_SNAPSHOTS = 1_000
_SNAPSHOT_SCOPE_INPUT_KEY = "__ag_ui_snapshot_scope"
_DEFAULT_STATE_INPUT_KEY = "__ag_ui_default_state"
@dataclass(slots=True)
class AGUIThreadSnapshot:
"""Replayable AG-UI Thread state.
AG-UI Thread Snapshots intentionally contain only data that can be replayed
to a UI: message snapshots, optional Shared State, and optional interruption
state. They do not include raw events, request metadata, auth claims,
diagnostics, traces, or provider responses.
Attributes:
messages: Replayable AG-UI message snapshots.
state: Optional AG-UI Shared State snapshot.
interrupt: Optional interruption state from ``RUN_FINISHED.interrupt``.
"""
messages: list[dict[str, Any]] = field(default_factory=list)
state: dict[str, Any] | None = None
interrupt: list[dict[str, Any]] | None = None
@runtime_checkable
class AGUIThreadSnapshotStore(Protocol):
"""Async store for latest AG-UI Thread Snapshots keyed by scope and thread id."""
async def save(
self,
*,
scope: SnapshotScope,
thread_id: AGUIThreadID,
snapshot: AGUIThreadSnapshot,
) -> None:
"""Save the latest snapshot for an AG-UI Thread within a Snapshot Scope.
Args:
scope: Application-defined Snapshot Scope. This is part of the
storage key and must represent the app's authorization boundary.
thread_id: AG-UI Thread id within the scope.
snapshot: Snapshot to save.
"""
...
async def get(
self,
*,
scope: SnapshotScope,
thread_id: AGUIThreadID,
) -> AGUIThreadSnapshot | None:
"""Get the latest snapshot for an AG-UI Thread within a Snapshot Scope.
Args:
scope: Application-defined Snapshot Scope.
thread_id: AG-UI Thread id within the scope.
Returns:
The latest snapshot, or ``None`` when no snapshot exists for the key.
"""
...
async def delete(
self,
*,
scope: SnapshotScope,
thread_id: AGUIThreadID,
) -> bool:
"""Delete the latest snapshot for an AG-UI Thread within a Snapshot Scope.
Args:
scope: Application-defined Snapshot Scope.
thread_id: AG-UI Thread id within the scope.
Returns:
``True`` when a snapshot was deleted, otherwise ``False``.
"""
...
async def clear(self, *, scope: SnapshotScope | None = None) -> None:
"""Clear saved snapshots.
Args:
scope: Optional Snapshot Scope to clear. When omitted, all in-memory
snapshots are cleared.
"""
...
class InMemoryAGUIThreadSnapshotStore:
"""Bounded memory-only latest snapshot store for local development, demos, and tests.
This store keeps at most one snapshot per ``(scope, thread_id)`` key. It is
process-local and not durable production storage.
"""
def __init__(self, *, max_snapshots: int = DEFAULT_MAX_THREAD_SNAPSHOTS) -> None:
"""Initialize the in-memory snapshot store.
Keyword Args:
max_snapshots: Maximum number of scoped thread snapshots to retain.
Raises:
ValueError: If ``max_snapshots`` is less than 1.
"""
if max_snapshots < 1:
raise ValueError("max_snapshots must be greater than 0.")
self._max_snapshots = max_snapshots
self._snapshots: dict[_SnapshotKey, AGUIThreadSnapshot] = {}
async def save(
self,
*,
scope: SnapshotScope,
thread_id: AGUIThreadID,
snapshot: AGUIThreadSnapshot,
) -> None:
"""Save the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
key = self._key(scope=scope, thread_id=thread_id)
if key in self._snapshots:
del self._snapshots[key]
self._snapshots[key] = copy.deepcopy(snapshot)
self._evict_oldest()
async def get(
self,
*,
scope: SnapshotScope,
thread_id: AGUIThreadID,
) -> AGUIThreadSnapshot | None:
"""Get the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
snapshot = self._snapshots.get(self._key(scope=scope, thread_id=thread_id))
return copy.deepcopy(snapshot) if snapshot is not None else None
async def delete(
self,
*,
scope: SnapshotScope,
thread_id: AGUIThreadID,
) -> bool:
"""Delete the latest snapshot for an AG-UI Thread within a Snapshot Scope."""
key = self._key(scope=scope, thread_id=thread_id)
if key not in self._snapshots:
return False
del self._snapshots[key]
return True
async def clear(self, *, scope: SnapshotScope | None = None) -> None:
"""Clear saved snapshots, optionally limited to one Snapshot Scope."""
if scope is None:
self._snapshots.clear()
return
normalized_scope = self._normalize_key_part(scope, "scope")
for key in list(self._snapshots):
if key[0] == normalized_scope:
del self._snapshots[key]
@classmethod
def _key(cls, *, scope: SnapshotScope, thread_id: AGUIThreadID) -> _SnapshotKey:
return (
cls._normalize_key_part(scope, "scope"),
cls._normalize_key_part(thread_id, "thread_id"),
)
@staticmethod
def _normalize_key_part(value: str, name: str) -> str:
if not isinstance(value, str):
raise TypeError(f"{name} must be a string.")
if not value:
raise ValueError(f"{name} must be a non-empty string.")
return value
def _evict_oldest(self) -> None:
while len(self._snapshots) > self._max_snapshots:
del self._snapshots[next(iter(self._snapshots))]
@@ -4,203 +4,18 @@
from __future__ import annotations
import copy
import logging
import uuid
from collections.abc import AsyncGenerator, Callable
from typing import Any, cast
from typing import Any
from ag_ui.core import (
BaseEvent,
MessagesSnapshotEvent,
RunErrorEvent,
RunFinishedEvent,
RunStartedEvent,
StateSnapshotEvent,
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from ag_ui.core import BaseEvent
from agent_framework import Workflow
from ._message_adapters import agui_messages_to_snapshot_format
from ._run_common import (
_build_run_finished_event,
_extract_resume_payload,
_reconstruct_messages_from_thread_snapshot,
)
from ._snapshots import (
_DEFAULT_STATE_INPUT_KEY,
_SNAPSHOT_SCOPE_INPUT_KEY,
AGUIThreadSnapshot,
AGUIThreadSnapshotStore,
)
from ._utils import generate_event_id, make_json_safe
from ._workflow_run import run_workflow_stream
logger = logging.getLogger(__name__)
WorkflowFactory = Callable[[str], Workflow]
def _event_messages_to_snapshot_dicts(messages: list[Any]) -> list[dict[str, Any]]:
"""Convert AG-UI message event models to plain snapshot dictionaries."""
safe_messages = make_json_safe(messages)
if not isinstance(safe_messages, list):
return []
return [cast(dict[str, Any], message) for message in safe_messages if isinstance(message, dict)]
class _WorkflowSnapshotBuilder:
"""Capture replayable workflow protocol output without retaining raw events."""
def __init__(self, raw_messages: list[dict[str, Any]]) -> None:
self._synthesized_messages = agui_messages_to_snapshot_format(raw_messages)
self._emitted_messages: list[dict[str, Any]] | None = None
self._open_text_message: dict[str, Any] | None = None
self._tool_call_message: dict[str, Any] | None = None
self._tool_calls_by_id: dict[str, dict[str, Any]] = {}
self.state: dict[str, Any] | None = None
self.interrupt: list[dict[str, Any]] | None = None
def observe(self, event: BaseEvent) -> None:
"""Fold one replayable AG-UI event into the latest snapshot state."""
if isinstance(event, StateSnapshotEvent):
state = make_json_safe(event.snapshot)
if isinstance(state, dict):
self.state = cast(dict[str, Any], state)
return
if isinstance(event, MessagesSnapshotEvent):
self._emitted_messages = _event_messages_to_snapshot_dicts(list(event.messages))
return
if isinstance(event, RunFinishedEvent):
interrupt = make_json_safe(getattr(event, "interrupt", None))
if isinstance(interrupt, list):
self.interrupt = [cast(dict[str, Any], item) for item in interrupt if isinstance(item, dict)]
return
if self._emitted_messages is not None:
return
if isinstance(event, TextMessageStartEvent):
self._observe_text_start(event)
elif isinstance(event, TextMessageContentEvent):
self._observe_text_content(event)
elif isinstance(event, TextMessageEndEvent):
self._observe_text_end(event)
elif isinstance(event, ToolCallStartEvent):
self._observe_tool_call_start(event)
elif isinstance(event, ToolCallArgsEvent):
self._observe_tool_call_args(event)
elif isinstance(event, ToolCallResultEvent):
self._observe_tool_call_result(event)
def build(self) -> AGUIThreadSnapshot:
"""Return the replayable thread snapshot."""
self._flush_open_text_message()
messages = self._emitted_messages if self._emitted_messages is not None else self._synthesized_messages
return AGUIThreadSnapshot(messages=messages, state=self.state, interrupt=self.interrupt)
def _observe_text_start(self, event: TextMessageStartEvent) -> None:
if self._open_text_message is not None and self._open_text_message.get("id") != event.message_id:
self._flush_open_text_message()
self._open_text_message = {"id": event.message_id, "role": event.role, "content": ""}
def _observe_text_content(self, event: TextMessageContentEvent) -> None:
if self._open_text_message is None or self._open_text_message.get("id") != event.message_id:
self._open_text_message = {"id": event.message_id, "role": "assistant", "content": ""}
self._open_text_message["content"] = f"{self._open_text_message.get('content', '')}{event.delta}"
def _observe_text_end(self, event: TextMessageEndEvent) -> None:
if self._open_text_message is None or self._open_text_message.get("id") != event.message_id:
return
self._flush_open_text_message()
def _observe_tool_call_start(self, event: ToolCallStartEvent) -> None:
parent_message_id = event.parent_message_id
if (
self._open_text_message is not None
and parent_message_id is not None
and self._open_text_message.get("id") == parent_message_id
and self._open_text_message.get("content")
):
self._open_text_message["id"] = generate_event_id()
self._flush_open_text_message()
if self._tool_call_message is None or (
parent_message_id is not None and self._tool_call_message.get("id") != parent_message_id
):
self._tool_call_message = {
"id": parent_message_id or generate_event_id(),
"role": "assistant",
"tool_calls": [],
}
self._synthesized_messages.append(self._tool_call_message)
tool_call = {
"id": event.tool_call_id,
"type": "function",
"function": {"name": event.tool_call_name, "arguments": ""},
}
cast(list[dict[str, Any]], self._tool_call_message["tool_calls"]).append(tool_call)
self._tool_calls_by_id[event.tool_call_id] = tool_call
def _observe_tool_call_args(self, event: ToolCallArgsEvent) -> None:
tool_call = self._tool_calls_by_id.get(event.tool_call_id)
if tool_call is None:
return
function_payload = cast(dict[str, Any], tool_call["function"])
function_payload["arguments"] = f"{function_payload.get('arguments', '')}{event.delta}"
def _observe_tool_call_result(self, event: ToolCallResultEvent) -> None:
self._synthesized_messages.append(
{
"id": event.message_id,
"role": "tool",
"toolCallId": event.tool_call_id,
"content": event.content,
}
)
# A result closes the current tool-call group; later tool calls start a new
# assistant message so replayed transcripts keep results adjacent to their
# tool_calls message, which provider APIs require.
self._tool_call_message = None
def _flush_open_text_message(self) -> None:
if self._open_text_message is None:
return
if self._open_text_message.get("content"):
self._synthesized_messages.append(self._open_text_message)
# Text between tool calls closes the current tool-call group as well.
self._tool_call_message = None
self._open_text_message = None
async def _hydrate_workflow_thread_snapshot(
*,
snapshot_store: AGUIThreadSnapshotStore,
scope: str,
thread_id: str,
run_id: str,
) -> AsyncGenerator[BaseEvent]:
"""Replay the latest stored workflow AG-UI Thread Snapshot without invoking the workflow."""
yield RunStartedEvent(run_id=run_id, thread_id=thread_id)
snapshot = await snapshot_store.get(scope=scope, thread_id=thread_id)
if snapshot is None:
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id)
return
if snapshot.state is not None:
yield StateSnapshotEvent(snapshot=snapshot.state)
if snapshot.messages:
yield MessagesSnapshotEvent(messages=snapshot.messages) # type: ignore[arg-type]
yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=snapshot.interrupt)
class AgentFrameworkWorkflow:
"""Base AG-UI workflow wrapper.
@@ -214,30 +29,15 @@ class AgentFrameworkWorkflow:
workflow_factory: WorkflowFactory | None = None,
name: str | None = None,
description: str | None = None,
snapshot_store: AGUIThreadSnapshotStore | None = None,
) -> None:
"""Initialize the AG-UI workflow wrapper.
Args:
workflow: Optional workflow instance to expose.
workflow_factory: Optional factory for thread-scoped workflow instances.
name: Optional workflow name.
description: Optional workflow description.
snapshot_store: Optional AG-UI Thread Snapshot store. Snapshot persistence remains inactive unless
endpoint setup also provides an explicit Snapshot Scope resolver.
"""
if workflow is not None and workflow_factory is not None:
raise ValueError("Pass either workflow= or workflow_factory=, not both.")
self.workflow = workflow
self._workflow_factory = workflow_factory
# Cache keyed by (snapshot_scope, thread_id): the Snapshot Scope is the
# authorization boundary, so the same thread id under different scopes
# must never share an in-memory workflow instance.
self._workflow_by_thread: dict[tuple[str | None, str], Workflow] = {}
self._workflow_by_thread: dict[str, Workflow] = {}
self.name = name if name is not None else getattr(workflow, "name", "workflow")
self.description = description if description is not None else getattr(workflow, "description", "")
self.snapshot_store = snapshot_store
@staticmethod
def _thread_id_from_input(input_data: dict[str, Any]) -> str:
@@ -247,7 +47,7 @@ class AgentFrameworkWorkflow:
return str(thread_id)
return str(uuid.uuid4())
def _resolve_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> Workflow:
def _resolve_workflow(self, thread_id: str) -> Workflow:
"""Get the workflow instance for the current run."""
if self.workflow is not None:
return self.workflow
@@ -255,22 +55,17 @@ class AgentFrameworkWorkflow:
if self._workflow_factory is None:
raise NotImplementedError("No workflow is attached. Override run or pass workflow=/workflow_factory=.")
cache_key = (snapshot_scope, thread_id)
workflow = self._workflow_by_thread.get(cache_key)
workflow = self._workflow_by_thread.get(thread_id)
if workflow is None:
workflow = self._workflow_factory(thread_id)
if not isinstance(workflow, Workflow):
raise TypeError("workflow_factory must return a Workflow instance.")
self._workflow_by_thread[cache_key] = workflow
self._workflow_by_thread[thread_id] = workflow
return workflow
def clear_thread_workflow(self, thread_id: str, snapshot_scope: str | None = None) -> None:
"""Drop cached workflow instances for a thread, optionally limited to one Snapshot Scope."""
if snapshot_scope is not None:
self._workflow_by_thread.pop((snapshot_scope, thread_id), None)
return
for key in [key for key in self._workflow_by_thread if key[1] == thread_id]:
del self._workflow_by_thread[key]
def clear_thread_workflow(self, thread_id: str) -> None:
"""Drop a single cached thread workflow instance."""
self._workflow_by_thread.pop(thread_id, None)
def clear_workflow_cache(self) -> None:
"""Drop all cached thread workflow instances."""
@@ -282,96 +77,6 @@ class AgentFrameworkWorkflow:
Subclasses may override this to provide custom AG-UI streams.
"""
thread_id = self._thread_id_from_input(input_data)
run_id = str(input_data.get("run_id") or input_data.get("runId") or uuid.uuid4())
snapshot_scope = cast(str | None, input_data.get(_SNAPSHOT_SCOPE_INPUT_KEY))
raw_messages = list(cast(list[dict[str, Any]], input_data.get("messages", []) or []))
resume_payload = _extract_resume_payload(input_data)
snapshot_store = self.snapshot_store
if snapshot_store is not None and snapshot_scope is not None and not raw_messages and resume_payload is None:
async for event in _hydrate_workflow_thread_snapshot(
snapshot_store=snapshot_store,
scope=snapshot_scope,
thread_id=thread_id,
run_id=run_id,
):
yield event
return
# Load the stored snapshot for follow-up turns so the workflow runs with the
# full persisted thread history instead of just the latest request messages.
stored_snapshot: AGUIThreadSnapshot | None = None
if snapshot_store is not None and snapshot_scope is not None:
stored_snapshot = await snapshot_store.get(scope=snapshot_scope, thread_id=thread_id)
if stored_snapshot is not None and resume_payload is None:
raw_messages = _reconstruct_messages_from_thread_snapshot(
stored_messages=stored_snapshot.messages,
incoming_messages=raw_messages,
stored_interrupt=stored_snapshot.interrupt,
)
input_data["messages"] = raw_messages
# Merge stored state with request overrides, then fill endpoint-deferred
# defaults only for keys missing from both.
request_state = input_data.get("state")
deferred_default_state = cast(dict[str, Any] | None, input_data.get(_DEFAULT_STATE_INPUT_KEY))
effective_state: dict[str, Any] = {}
if stored_snapshot is not None and stored_snapshot.state is not None:
effective_state.update(stored_snapshot.state)
if isinstance(request_state, dict):
effective_state.update(cast(dict[str, Any], request_state))
if deferred_default_state:
for key, value in deferred_default_state.items():
if key not in effective_state:
effective_state[key] = copy.deepcopy(value)
if effective_state:
input_data["state"] = effective_state
workflow = self._resolve_workflow(thread_id, snapshot_scope)
builder_seed_messages = raw_messages
if resume_payload is not None and stored_snapshot is not None:
# Resume requests carry only the synthesized interrupt response, so seed
# the builder with stored history to avoid persisting a truncated thread.
builder_seed_messages = [
copy.deepcopy(message) for message in stored_snapshot.messages
] + builder_seed_messages
snapshot_builder = (
_WorkflowSnapshotBuilder(builder_seed_messages)
if snapshot_store is not None and snapshot_scope is not None
else None
)
if snapshot_builder is not None and effective_state:
# Seed builder state so a run that emits no StateSnapshotEvent still
# persists the latest known Shared State instead of dropping it.
state_snapshot = make_json_safe(effective_state)
if isinstance(state_snapshot, dict):
snapshot_builder.state = cast(dict[str, Any], state_snapshot)
run_error_emitted = False
workflow = self._resolve_workflow(thread_id)
async for event in run_workflow_stream(input_data, workflow):
if snapshot_builder is not None:
snapshot_builder.observe(event)
if isinstance(event, RunErrorEvent):
run_error_emitted = True
yield event
if (
snapshot_builder is not None
and not run_error_emitted
and snapshot_store is not None
and snapshot_scope is not None
):
try:
await snapshot_store.save(
scope=snapshot_scope,
thread_id=thread_id,
snapshot=snapshot_builder.build(),
)
except Exception:
# RUN_FINISHED has already been yielded; a store failure must not
# surface as a second terminal RUN_ERROR event. The previous
# snapshot stays available for hydration.
logger.exception(
"Failed to save AG-UI Thread Snapshot for scope=%s thread_id=%s; keeping previous snapshot.",
snapshot_scope,
thread_id,
)
File diff suppressed because it is too large Load Diff
@@ -32,21 +32,6 @@ def test_agent_framework_ag_ui_exports_state_update() -> None:
assert callable(state_update)
def test_agent_framework_ag_ui_exports_snapshot_primitives() -> None:
"""Runtime package should export AG-UI Thread Snapshot primitives."""
from agent_framework_ag_ui import (
DEFAULT_MAX_THREAD_SNAPSHOTS,
AGUIThreadSnapshot,
AGUIThreadSnapshotStore,
InMemoryAGUIThreadSnapshotStore,
)
assert AGUIThreadSnapshot.__name__ == "AGUIThreadSnapshot"
assert AGUIThreadSnapshotStore.__name__ == "AGUIThreadSnapshotStore"
assert InMemoryAGUIThreadSnapshotStore.__name__ == "InMemoryAGUIThreadSnapshotStore"
assert DEFAULT_MAX_THREAD_SNAPSHOTS >= 1
def test_core_ag_ui_lazy_exports_include_event_converter_and_http_service() -> None:
"""Core facade must expose AGUIEventConverter, AGUIHttpService, and __version__."""
from agent_framework import ag_ui
@@ -54,13 +39,3 @@ def test_core_ag_ui_lazy_exports_include_event_converter_and_http_service() -> N
assert hasattr(ag_ui, "AGUIEventConverter")
assert hasattr(ag_ui, "AGUIHttpService")
assert hasattr(ag_ui, "__version__")
def test_core_ag_ui_lazy_exports_include_snapshot_primitives() -> None:
"""Core facade must expose snapshot primitives needed for endpoint configuration."""
from agent_framework import ag_ui
assert hasattr(ag_ui, "AGUIThreadSnapshot")
assert hasattr(ag_ui, "AGUIThreadSnapshotStore")
assert hasattr(ag_ui, "InMemoryAGUIThreadSnapshotStore")
assert hasattr(ag_ui, "SnapshotScopeResolver")
@@ -1,160 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for AG-UI thread snapshot storage primitives."""
from dataclasses import fields
from agent_framework_ag_ui import AGUIThreadSnapshot, AGUIThreadSnapshotStore, InMemoryAGUIThreadSnapshotStore
def test_thread_snapshot_model_contains_only_replayable_snapshot_fields() -> None:
"""The public snapshot model is limited to messages, Shared State, and interruption state."""
assert [field.name for field in fields(AGUIThreadSnapshot)] == ["messages", "state", "interrupt"]
def test_in_memory_snapshot_store_satisfies_snapshot_store_protocol() -> None:
"""The built-in store conforms to the public async store protocol."""
assert isinstance(InMemoryAGUIThreadSnapshotStore(), AGUIThreadSnapshotStore)
async def test_in_memory_snapshot_store_replaces_latest_snapshot() -> None:
"""Saving the same scoped thread key replaces the previous snapshot."""
store = InMemoryAGUIThreadSnapshotStore()
await store.save(
scope="tenant-a",
thread_id="thread-1",
snapshot=AGUIThreadSnapshot(messages=[{"id": "first"}], state={"count": 1}),
)
await store.save(
scope="tenant-a",
thread_id="thread-1",
snapshot=AGUIThreadSnapshot(messages=[{"id": "second"}], state={"count": 2}),
)
snapshot = await store.get(scope="tenant-a", thread_id="thread-1")
assert snapshot is not None
assert snapshot.messages == [{"id": "second"}]
assert snapshot.state == {"count": 2}
async def test_in_memory_snapshot_store_keeps_scopes_separate() -> None:
"""The same AG-UI Thread id in different Snapshot Scopes addresses different snapshots."""
store = InMemoryAGUIThreadSnapshotStore()
await store.save(
scope="tenant-a",
thread_id="thread-1",
snapshot=AGUIThreadSnapshot(messages=[{"id": "a", "role": "user", "content": "from a"}]),
)
await store.save(
scope="tenant-b",
thread_id="thread-1",
snapshot=AGUIThreadSnapshot(messages=[{"id": "b", "role": "user", "content": "from b"}]),
)
tenant_a_snapshot = await store.get(scope="tenant-a", thread_id="thread-1")
tenant_b_snapshot = await store.get(scope="tenant-b", thread_id="thread-1")
assert tenant_a_snapshot is not None
assert tenant_b_snapshot is not None
assert tenant_a_snapshot.messages == [{"id": "a", "role": "user", "content": "from a"}]
assert tenant_b_snapshot.messages == [{"id": "b", "role": "user", "content": "from b"}]
async def test_in_memory_snapshot_store_deletes_and_clears_snapshots() -> None:
"""Delete removes one scoped thread key, while clear can remove a scope or the whole store."""
store = InMemoryAGUIThreadSnapshotStore()
await store.save(scope="tenant-a", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "a1"}]))
await store.save(scope="tenant-a", thread_id="thread-2", snapshot=AGUIThreadSnapshot(messages=[{"id": "a2"}]))
await store.save(scope="tenant-b", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "b1"}]))
assert await store.delete(scope="tenant-a", thread_id="thread-1") is True
assert await store.delete(scope="tenant-a", thread_id="thread-1") is False
assert await store.get(scope="tenant-a", thread_id="thread-1") is None
assert await store.get(scope="tenant-a", thread_id="thread-2") is not None
await store.clear(scope="tenant-a")
assert await store.get(scope="tenant-a", thread_id="thread-2") is None
assert await store.get(scope="tenant-b", thread_id="thread-1") is not None
await store.clear()
assert await store.get(scope="tenant-b", thread_id="thread-1") is None
async def test_in_memory_snapshot_store_evicts_oldest_snapshot_when_bounded() -> None:
"""The memory store bounds retained scoped thread snapshots."""
store = InMemoryAGUIThreadSnapshotStore(max_snapshots=2)
await store.save(scope="tenant-a", thread_id="thread-1", snapshot=AGUIThreadSnapshot(messages=[{"id": "first"}]))
await store.save(scope="tenant-a", thread_id="thread-2", snapshot=AGUIThreadSnapshot(messages=[{"id": "second"}]))
await store.save(scope="tenant-a", thread_id="thread-3", snapshot=AGUIThreadSnapshot(messages=[{"id": "third"}]))
assert await store.get(scope="tenant-a", thread_id="thread-1") is None
assert await store.get(scope="tenant-a", thread_id="thread-2") is not None
assert await store.get(scope="tenant-a", thread_id="thread-3") is not None
def test_workflow_snapshot_builder_splits_tool_call_groups() -> None:
"""Tool calls separated by results or text synthesize provider-valid message groups."""
from ag_ui.core import (
TextMessageContentEvent,
TextMessageEndEvent,
TextMessageStartEvent,
ToolCallArgsEvent,
ToolCallResultEvent,
ToolCallStartEvent,
)
from agent_framework_ag_ui._workflow import _WorkflowSnapshotBuilder
builder = _WorkflowSnapshotBuilder([])
builder.observe(ToolCallStartEvent(tool_call_id="call-a", tool_call_name="toolA"))
builder.observe(ToolCallArgsEvent(tool_call_id="call-a", delta='{"x": 1}'))
builder.observe(ToolCallResultEvent(message_id="result-a", tool_call_id="call-a", content="resA"))
builder.observe(TextMessageStartEvent(message_id="text-1", role="assistant"))
builder.observe(TextMessageContentEvent(message_id="text-1", delta="thinking"))
builder.observe(TextMessageEndEvent(message_id="text-1"))
builder.observe(ToolCallStartEvent(tool_call_id="call-b", tool_call_name="toolB"))
builder.observe(ToolCallResultEvent(message_id="result-b", tool_call_id="call-b", content="resB"))
messages = builder.build().messages
shapes = [
(
message.get("role"),
[tool_call["id"] for tool_call in message.get("tool_calls", [])] or message.get("toolCallId"),
)
for message in messages
]
assert shapes == [
("assistant", ["call-a"]),
("tool", "call-a"),
("assistant", None),
("assistant", ["call-b"]),
("tool", "call-b"),
]
async def test_in_memory_snapshot_store_rejects_invalid_keys() -> None:
"""Key parts must be non-empty strings for every store operation."""
import pytest
store = InMemoryAGUIThreadSnapshotStore()
snapshot = AGUIThreadSnapshot()
with pytest.raises(ValueError):
await store.save(scope="", thread_id="thread-1", snapshot=snapshot)
with pytest.raises(ValueError):
await store.save(scope="tenant-a", thread_id="", snapshot=snapshot)
with pytest.raises(TypeError):
await store.save(scope=123, thread_id="thread-1", snapshot=snapshot) # type: ignore[arg-type]
with pytest.raises(ValueError):
await store.get(scope="tenant-a", thread_id="")
with pytest.raises(TypeError):
await store.delete(scope=None, thread_id="thread-1") # type: ignore[arg-type]
with pytest.raises(ValueError):
await store.clear(scope="")
@@ -1024,10 +1024,8 @@ class RawAnthropicClient(
usage_details["input_token_count"] = usage.input_tokens
if usage.cache_creation_input_tokens is not None:
usage_details["anthropic.cache_creation_input_tokens"] = usage.cache_creation_input_tokens # type: ignore[typeddict-unknown-key]
usage_details["cache_creation_input_token_count"] = usage.cache_creation_input_tokens
if usage.cache_read_input_tokens is not None:
usage_details["anthropic.cache_read_input_tokens"] = usage.cache_read_input_tokens # type: ignore[typeddict-unknown-key]
usage_details["cache_read_input_token_count"] = usage.cache_read_input_tokens
return usage_details
def _parse_contents_from_anthropic(
@@ -2354,27 +2354,6 @@ def test_parse_usage_with_cache_tokens(mock_anthropic_client: MagicMock) -> None
assert result["input_token_count"] == 100
assert result["anthropic.cache_creation_input_tokens"] == 20
assert result["anthropic.cache_read_input_tokens"] == 30
assert result["cache_creation_input_token_count"] == 20
assert result["cache_read_input_token_count"] == 30
def test_parse_usage_preserves_zero_cache_tokens(mock_anthropic_client: MagicMock) -> None:
"""Test parsing usage preserves zero-valued mapped cache tokens."""
client = create_test_anthropic_client(mock_anthropic_client)
mock_usage = MagicMock()
mock_usage.input_tokens = 100
mock_usage.output_tokens = 50
mock_usage.cache_creation_input_tokens = 0
mock_usage.cache_read_input_tokens = 0
result = client._parse_usage_from_anthropic(mock_usage)
assert result is not None
assert result["anthropic.cache_creation_input_tokens"] == 0
assert result["cache_creation_input_token_count"] == 0
assert result["anthropic.cache_read_input_tokens"] == 0
assert result["cache_read_input_token_count"] == 0
# Code Execution Result Tests
+2 -7
View File
@@ -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`, `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.
- **`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.
- **`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_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.
- **`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.
### Tool Approval Harness (`_harness/_tool_approval.py`)
@@ -116,11 +116,6 @@ 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/`)
@@ -102,12 +102,6 @@ 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,
@@ -131,6 +125,7 @@ from ._harness._todo import (
TodoSessionStore,
TodoStore,
)
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPTaskOptions, MCPWebsocketTool, SamplingApprovalCallback
from ._harness._tool_approval import (
DEFAULT_TOOL_APPROVAL_SOURCE_ID,
ToolApprovalMiddleware,
@@ -140,7 +135,6 @@ from ._harness._tool_approval import (
create_always_approve_tool_response,
create_always_approve_tool_with_arguments_response,
)
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPTaskOptions, MCPWebsocketTool, SamplingApprovalCallback
from ._middleware import (
AgentContext,
AgentMiddleware,
@@ -369,7 +363,6 @@ __all__ = [
"AgentExecutorResponse",
"AgentFileStore",
"AgentFrameworkException",
"AgentLoopMiddleware",
"AgentMiddleware",
"AgentMiddlewareLayer",
"AgentMiddlewareTypes",
@@ -461,7 +454,6 @@ __all__ = [
"InlineSkill",
"InlineSkillResource",
"InlineSkillScript",
"JudgeVerdict",
"LocalEvaluator",
"MCPSkill",
"MCPSkillResource",
@@ -566,7 +558,6 @@ __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",
@@ -597,7 +588,6 @@ __all__ = [
"response_handler",
"set_agent_mode",
"step",
"todos_remaining",
"tool",
"tool_call_args_match",
"tool_called_check",
@@ -5,10 +5,9 @@
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 six tools — ``file_access_save_file``,
sessions and agents. The provider exposes five tools — ``file_access_save_file``,
``file_access_read_file``, ``file_access_delete_file``, ``file_access_list_files``,
``file_access_list_subdirectories``, and ``file_access_search_files`` — by
registering them on the per-invocation
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
@@ -49,11 +48,7 @@ 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.\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."
"asked you to do so."
)
# Maximum number of characters of context to include on either side of the first
@@ -183,16 +178,10 @@ 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).
``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``.
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).
"""
if pattern is None or not pattern.strip():
return True
@@ -429,18 +418,6 @@ 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``.
@@ -455,8 +432,6 @@ 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``.
@@ -466,19 +441,12 @@ 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. 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.
filter which files are searched. When ``None`` or blank, every
file in the directory is searched.
Returns:
The list of files whose content matched, with snippet and matching
line metadata. Each result's ``file_name`` is the path relative to
``directory`` (forward slashes).
line metadata.
"""
@abstractmethod
@@ -562,38 +530,6 @@ 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)
@@ -605,8 +541,6 @@ 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.
@@ -614,10 +548,7 @@ 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. 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.
what :class:`FileSystemAgentFileStore` would produce.
"""
prefix = _normalize_relative_path(directory, is_directory=True).lower()
if prefix and not prefix.endswith("/"):
@@ -633,7 +564,7 @@ class InMemoryAgentFileStore(AgentFileStore):
if not key.startswith(prefix):
continue
relative_key = key[len(prefix) :]
if not recursive and "/" in relative_key:
if "/" in relative_key:
continue
relative_display = display[len(prefix) :]
if not _matches_glob(relative_display, file_pattern):
@@ -864,28 +795,6 @@ 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)
@@ -900,8 +809,6 @@ 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.
@@ -909,50 +816,23 @@ 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". 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.
not searchable".
"""
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, recursive))
return await _run_search_with_timeout(lambda: self._search_files_sync(full_dir, regex, file_pattern))
@staticmethod
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]:
def _search_files_sync(full_dir: Path, regex: re.Pattern[str], file_pattern: str | None) -> list[FileSearchResult]:
if not full_dir.is_dir():
return []
results: list[FileSearchResult] = []
skipped: list[str] = []
for relative_name, entry in FileSystemAgentFileStore._enumerate_search_files(full_dir, recursive):
if not _matches_glob(relative_name, file_pattern):
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):
continue
try:
file_content = entry.read_text(encoding="utf-8")
@@ -961,9 +841,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(relative_name)
skipped.append(file_name)
continue
result = _search_file_content(relative_name, file_content, regex)
result = _search_file_content(file_name, file_content, regex)
if result is not None:
results.append(result)
if skipped:
@@ -985,18 +865,15 @@ class FileSystemAgentFileStore(AgentFileStore):
class FileAccessProvider(ContextProvider):
"""Context provider that gives an agent CRUD/search access to a shared file store.
The provider exposes six tools to the agent via the per-invocation
The provider exposes five 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 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.
- ``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.
Unlike :class:`~agent_framework.MemoryContextProvider`, which provides
session-scoped memory that may be isolated per session,
@@ -1099,45 +976,17 @@ 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 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.
"""
"""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
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("", regex_pattern, pattern, recursive=True)
results = await self.store.search_files(target, regex_pattern, pattern)
except ValueError as exc:
return f"Could not search files: {exc}"
except OSError as exc:
@@ -1152,7 +1001,6 @@ class FileAccessProvider(ContextProvider):
file_access_read_file,
file_access_delete_file,
file_access_list_files,
file_access_list_subdirectories,
file_access_search_files,
],
)
@@ -1,796 +0,0 @@
# 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
@@ -1989,7 +1989,9 @@ def _store_already_approved_approval_requests(
return
existing_groups = state.get(_ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY)
pending_groups: list[Any] = list(cast(Iterable[Any], existing_groups)) if isinstance(existing_groups, list) else []
pending_groups: list[Any] = (
list(cast(Iterable[Any], existing_groups)) if isinstance(existing_groups, list) else []
)
pending_groups.append({
"approval_request_ids": visible_ids,
"approval_requests": [request.to_dict() for request in already_approved_requests],
@@ -400,18 +400,12 @@ 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:
@@ -11,10 +11,6 @@ Supported classes and functions:
- AGUIChatClient
- AGUIEventConverter
- AGUIHttpService
- AGUIThreadSnapshot
- AGUIThreadSnapshotStore
- InMemoryAGUIThreadSnapshotStore
- SnapshotScopeResolver
- add_agent_framework_fastapi_endpoint
- state_update
- __version__
@@ -32,10 +28,6 @@ _IMPORTS = [
"AGUIChatClient",
"AGUIEventConverter",
"AGUIHttpService",
"AGUIThreadSnapshot",
"AGUIThreadSnapshotStore",
"InMemoryAGUIThreadSnapshotStore",
"SnapshotScopeResolver",
"state_update",
"__version__",
]
@@ -6,10 +6,6 @@ from agent_framework_ag_ui import (
AGUIChatClient,
AGUIEventConverter,
AGUIHttpService,
AGUIThreadSnapshot,
AGUIThreadSnapshotStore,
InMemoryAGUIThreadSnapshotStore,
SnapshotScopeResolver,
__version__,
add_agent_framework_fastapi_endpoint,
state_update,
@@ -19,12 +15,8 @@ __all__ = [
"AGUIChatClient",
"AGUIEventConverter",
"AGUIHttpService",
"AGUIThreadSnapshot",
"AGUIThreadSnapshotStore",
"AgentFrameworkAgent",
"AgentFrameworkWorkflow",
"InMemoryAGUIThreadSnapshotStore",
"SnapshotScopeResolver",
"__version__",
"add_agent_framework_fastapi_endpoint",
"state_update",
@@ -201,9 +201,6 @@ 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"
@@ -330,20 +327,6 @@ 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
@@ -2367,16 +2350,12 @@ def _apply_accumulated_usage(attributes: dict[str, Any], captured_fields: set[st
accumulated = INNER_ACCUMULATED_USAGE.get()
if not accumulated:
return
_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)
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
def _get_response_attributes(
@@ -2399,7 +2378,12 @@ def _get_response_attributes(
if model := getattr(response, "model", None):
attributes[OtelAttr.RESPONSE_MODEL] = model
if capture_usage and (usage := response.usage_details):
_apply_usage_attributes(attributes, usage)
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
return attributes
@@ -2423,9 +2407,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)) is not None:
if token_usage_histogram and (input_tokens := attributes.get(OtelAttr.INPUT_TOKENS)):
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)) is not None:
if token_usage_histogram and (output_tokens := attributes.get(OtelAttr.OUTPUT_TOKENS)):
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:
@@ -158,69 +158,6 @@ 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()
@@ -330,78 +267,6 @@ 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)
@@ -438,7 +303,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 six tools."""
"""``FileAccessProvider.before_run`` should add the canonical instructions and five tools."""
session = AgentSession(session_id="session-1")
store = InMemoryAgentFileStore()
provider = FileAccessProvider(store=store)
@@ -456,7 +321,6 @@ 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
@@ -490,7 +354,6 @@ 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"
@@ -533,7 +396,6 @@ 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"})
@@ -564,15 +426,6 @@ 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
@@ -581,12 +434,14 @@ 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 is recursive from the store root; scope to a subtree using
# the glob (``*`` crosses ``/`` with fnmatch). Results use root-relative names.
# The search tool should likewise accept an optional directory argument so
# agents can scope a search to a subfolder.
await save_file.invoke(arguments={"file_name": "reports/issues.md", "content": "ERROR nested"})
scoped = await search_files.invoke(arguments={"regex_pattern": "error", "file_pattern": "reports/*"})
scoped = await search_files.invoke(
arguments={"regex_pattern": "error", "file_pattern": "*.md", "directory": "reports"}
)
scoped_parsed = json.loads(scoped[0].text)
assert [entry["file_name"] for entry in scoped_parsed] == ["reports/issues.md"]
assert [entry["file_name"] for entry in scoped_parsed] == ["issues.md"]
deleted = await delete_file.invoke(arguments={"file_name": "plan.md"})
assert "deleted" in deleted[0].text
File diff suppressed because it is too large Load Diff
@@ -2154,58 +2154,6 @@ 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
@@ -2216,22 +2164,13 @@ 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,
"cache_creation_input_token_count": 10,
"cache_read_input_token_count": 20,
"reasoning_output_token_count": 30,
}
response.usage_details = {"input_token_count": 100, "output_token_count": 50}
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():
@@ -2994,23 +2933,6 @@ 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.
@@ -4015,21 +3937,11 @@ 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,
cache_read_input_token_count=100,
reasoning_output_token_count=25,
),
usage_details=UsageDetails(input_token_count=2239, output_token_count=192),
),
ChatResponse(
messages=Message(role="assistant", contents=["The weather in Seattle is sunny."]),
usage_details=UsageDetails(
input_token_count=2569,
output_token_count=99,
cache_read_input_token_count=200,
reasoning_output_token_count=0,
),
usage_details=UsageDetails(input_token_count=2569, output_token_count=99),
),
]
@@ -4053,18 +3965,12 @@ 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)
@@ -2979,16 +2979,10 @@ class RawOpenAIChatClient( # type: ignore[misc]
output_token_count=usage.output_tokens,
total_token_count=usage.total_tokens,
)
if usage.input_tokens_details:
cached_tokens = cast("int | None", getattr(usage.input_tokens_details, "cached_tokens", None))
if cached_tokens is not None:
details["openai.cached_input_tokens"] = cached_tokens # type: ignore[typeddict-unknown-key]
details["cache_read_input_token_count"] = cached_tokens
if usage.output_tokens_details:
reasoning_tokens = cast("int | None", getattr(usage.output_tokens_details, "reasoning_tokens", None))
if reasoning_tokens is not None:
details["openai.reasoning_tokens"] = reasoning_tokens # type: ignore[typeddict-unknown-key]
details["reasoning_output_token_count"] = reasoning_tokens
if usage.input_tokens_details and usage.input_tokens_details.cached_tokens:
details["openai.cached_input_tokens"] = usage.input_tokens_details.cached_tokens # type: ignore[typeddict-unknown-key]
if usage.output_tokens_details and usage.output_tokens_details.reasoning_tokens:
details["openai.reasoning_tokens"] = usage.output_tokens_details.reasoning_tokens # type: ignore[typeddict-unknown-key]
return details
def _get_metadata_from_response(self, output: Any) -> dict[str, Any]:
@@ -765,17 +765,15 @@ class RawOpenAIChatCompletionClient( # type: ignore[misc]
details["completion/accepted_prediction_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if tokens := usage.completion_tokens_details.audio_tokens:
details["completion/audio_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if (tokens := usage.completion_tokens_details.reasoning_tokens) is not None:
if tokens := usage.completion_tokens_details.reasoning_tokens:
details["completion/reasoning_tokens"] = tokens # type: ignore[typeddict-unknown-key]
details["reasoning_output_token_count"] = tokens
if tokens := usage.completion_tokens_details.rejected_prediction_tokens:
details["completion/rejected_prediction_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if usage.prompt_tokens_details:
if tokens := usage.prompt_tokens_details.audio_tokens:
details["prompt/audio_tokens"] = tokens # type: ignore[typeddict-unknown-key]
if (tokens := usage.prompt_tokens_details.cached_tokens) is not None:
if tokens := usage.prompt_tokens_details.cached_tokens:
details["prompt/cached_tokens"] = tokens # type: ignore[typeddict-unknown-key]
details["cache_read_input_token_count"] = tokens
return details
def _parse_text_from_openai(self, choice: Choice | ChunkChoice) -> Content | None:
@@ -3301,7 +3301,6 @@ def test_usage_details_with_cached_tokens() -> None:
assert details is not None
assert details["input_token_count"] == 200
assert details["openai.cached_input_tokens"] == 25
assert details["cache_read_input_token_count"] == 25
def test_usage_details_with_reasoning_tokens() -> None:
@@ -3320,49 +3319,6 @@ def test_usage_details_with_reasoning_tokens() -> None:
assert details is not None
assert details["output_token_count"] == 80
assert details["openai.reasoning_tokens"] == 30
assert details["reasoning_output_token_count"] == 30
def test_usage_details_with_zero_cached_and_reasoning_tokens() -> None:
"""Test _parse_usage_from_openai preserves zero-valued mapped usage details."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
mock_usage = MagicMock()
mock_usage.input_tokens = 150
mock_usage.output_tokens = 80
mock_usage.total_tokens = 230
mock_usage.input_tokens_details = MagicMock()
mock_usage.input_tokens_details.cached_tokens = 0
mock_usage.output_tokens_details = MagicMock()
mock_usage.output_tokens_details.reasoning_tokens = 0
details = client._parse_usage_from_openai(mock_usage) # type: ignore
assert details is not None
assert details["openai.cached_input_tokens"] == 0
assert details["cache_read_input_token_count"] == 0
assert details["openai.reasoning_tokens"] == 0
assert details["reasoning_output_token_count"] == 0
def test_usage_details_omits_missing_cached_and_reasoning_tokens() -> None:
"""Test _parse_usage_from_openai omits missing mapped usage details."""
client = OpenAIChatClient(model="test-model", api_key="test-key")
mock_usage = MagicMock()
mock_usage.input_tokens = 150
mock_usage.output_tokens = 80
mock_usage.total_tokens = 230
mock_usage.input_tokens_details = MagicMock()
mock_usage.input_tokens_details.cached_tokens = None
mock_usage.output_tokens_details = MagicMock()
mock_usage.output_tokens_details.reasoning_tokens = None
details = client._parse_usage_from_openai(mock_usage) # type: ignore
assert details is not None
assert "openai.cached_input_tokens" not in details
assert "cache_read_input_token_count" not in details
assert "openai.reasoning_tokens" not in details
assert "reasoning_output_token_count" not in details
def test_get_metadata_from_response() -> None:
@@ -1099,31 +1099,6 @@ def test_usage_content_in_streaming_response(
assert usage_content.usage_details["total_token_count"] == 150
def test_parse_usage_includes_standard_and_legacy_mapped_token_details() -> None:
"""Test _parse_usage_from_openai emits standard and legacy mapped token details."""
client = OpenAIChatCompletionClient(model="test-model", api_key="test-key")
mock_usage = MagicMock()
mock_usage.prompt_tokens = 100
mock_usage.completion_tokens = 50
mock_usage.total_tokens = 150
mock_usage.completion_tokens_details = MagicMock()
mock_usage.completion_tokens_details.accepted_prediction_tokens = None
mock_usage.completion_tokens_details.audio_tokens = None
mock_usage.completion_tokens_details.reasoning_tokens = 0
mock_usage.completion_tokens_details.rejected_prediction_tokens = None
mock_usage.prompt_tokens_details = MagicMock()
mock_usage.prompt_tokens_details.audio_tokens = None
mock_usage.prompt_tokens_details.cached_tokens = 0
details = client._parse_usage_from_openai(mock_usage) # type: ignore[arg-type]
assert details["completion/reasoning_tokens"] == 0
assert details["reasoning_output_token_count"] == 0
assert details["prompt/cached_tokens"] == 0
assert details["cache_read_input_token_count"] == 0
def test_streaming_chunk_with_usage_and_text(
openai_unit_test_env: dict[str, str],
) -> None:
@@ -38,9 +38,7 @@ the file_access_* tools.
## Getting started
- Start by listing available files with file_access_list_files to see what data
is available. Files may be organized into subdirectories — use
file_access_list_subdirectories to discover folders and explore the tree level
by level.
is available.
- Read the files to understand their structure and contents.
## Working with data
@@ -88,7 +86,7 @@ async def main() -> None:
# 3. Wire up the file access provider against a file-system-backed store
# rooted at the sample's working/ folder. The provider injects its
# default instructions plus exposes six file_access_* tools to the
# default instructions plus exposes five file_access_* tools to the
# agent for the duration of each run.
file_access = FileAccessProvider(store=FileSystemAgentFileStore(working_dir))
@@ -41,17 +41,7 @@ class PlanningQuestion(BaseModel):
choices: list[str] | None = Field(
default=None,
description=(
"For clarifications, this has a list of options that the user can "
"choose from. null for approvals.\n\n"
"Note: for clarifications, the user will always also be presented with "
"a free form input option, so make sure that each choice provided here "
"is a valid input for the next turn.\n"
'E.g. if the question is "Which stock are you referring to?" then valid '
'choices might be ["AAPL", "MSFT", "GOOG"], and the user could also type '
"their own answer.\n"
'Invalid choices would be ["Enter tickers directly", "Paste tickers"], '
"since these conflict with the already existing freeform option, and "
"don't directly provide valid inputs for the next turn."
"For clarifications, this has a list of options that the user can choose from. null for approvals."
),
)
@@ -7,10 +7,6 @@ This folder contains focused middleware samples for `Agent`, chat clients, tools
| File | Description |
|------|-------------|
| [`agent_and_run_level_middleware.py`](./agent_and_run_level_middleware.py) | Demonstrates combining agent-level and run-level middleware. |
| [`agent_loop_middleware_refinement.py`](./agent_loop_middleware_refinement.py) | Demonstrates `AgentLoopMiddleware` with a `should_continue` predicate: a completion-marker refinement loop with feedback tracking and `fresh_context`. |
| [`agent_loop_middleware_todos.py`](./agent_loop_middleware_todos.py) | Demonstrates `AgentLoopMiddleware` with a `should_continue` predicate built from a `TodoProvider` via `todos_remaining`, so the agent keeps working while open todos remain. |
| [`agent_loop_middleware_judge.py`](./agent_loop_middleware_judge.py) | Demonstrates `AgentLoopMiddleware.with_judge`: a ChatClient judge re-runs the agent until it decides the original request was answered, with `criteria` shared between the agent and the judge. |
| [`agent_loop_middleware_report.py`](./agent_loop_middleware_report.py) | Demonstrates composing two `AgentLoopMiddleware` on one agent: an inner `todos_remaining` loop that drafts a report todo-by-todo, wrapped by an outer report-style `with_judge` loop that re-runs it until an editor chat client judges the report publication-ready. |
| [`chat_middleware.py`](./chat_middleware.py) | Shows class-based and function-based chat middleware that can observe, modify, and override model calls. |
| [`class_based_middleware.py`](./class_based_middleware.py) | Shows class-based agent and function middleware. |
| [`decorator_middleware.py`](./decorator_middleware.py) | Demonstrates middleware registration with decorators. |
@@ -1,118 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import Agent, AgentLoopMiddleware
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
Agent Loop Middleware: ChatClient judge
This sample demonstrates ``AgentLoopMiddleware.with_judge(...)``: a second chat client decides (via a
``JudgeVerdict`` structured output) whether the original request was answered, and the loop continues
while the answer is "no". The judge's ``reasoning`` is fed back to the agent as the next iteration's
input, so the agent knows what is missing. The loop also passes a list of ``criteria``, which are
injected as an extra instruction for the agent and rendered into the judge's instructions.
The loop is run with streaming, so the judge's feedback between iterations shows up as a ``user``
update; the stream is printed as ``<role>: <content>`` lines.
Environment variables:
FOUNDRY_PROJECT_ENDPOINT — Azure AI Foundry project endpoint URL
FOUNDRY_MODEL — Model deployment name
Authentication:
Run ``az login`` before running this sample.
"""
async def judge_loop(client: FoundryChatClient, judge_client: FoundryChatClient) -> None:
"""A second chat client judges whether the request was answered."""
print("\n=== ChatClient judge (loop until the request is answered) ===")
# 1. Provide a ``judge_client``. The middleware asks it (via a ``JudgeVerdict`` structured
# output) whether the original request has been fully addressed and continues while the
# answer is "no". The judge's ``reasoning`` is fed back to the agent as the next iteration's
# input, so the agent knows what is missing. Judge loops default to a small ``max_iterations``
# cap because each pass costs an extra model call.
#
# ``criteria`` is a list of requirements the response must satisfy. The loop (a) injects them
# as an extra instruction for the agent before it runs and (b) renders them into the judge's
# instructions (the default judge prompt includes a ``{{criteria}}`` placeholder). Supply your
# own ``instructions`` string with ``{{criteria}}`` to control the wording, or omit ``criteria``
# entirely and pass a plain ``instructions`` string.
loop = AgentLoopMiddleware.with_judge(
judge_client,
criteria=[
"Mentions the moon",
"Includes at least one good joke",
"Is written as a single piece of fluent prose",
],
max_iterations=4,
)
agent = Agent(
client=client,
name="answerer",
instructions="You are a helpful assistant. Answer the user's question thoroughly.",
middleware=[loop],
)
# 2. Run with streaming; the judge's feedback appears as a ``user`` update between iterations
# until the judge is satisfied (or the iteration cap is reached). Each contiguous ``user``
# block marks the boundary into the next iteration, so we count loop iterations by those
# boundaries (robust to function calling, where one iteration may issue several model calls).
iterations = 1
in_user_block = False
assistant_open = False
async for update in agent.run("Explain why the sky is blue and sunsets are red.", stream=True):
if update.role == "user":
if not in_user_block:
iterations += 1
in_user_block = True
assistant_open = False
print(f"\nuser: {update.text}", flush=True)
continue
in_user_block = False
if update.text:
if not assistant_open:
print("\nassistant: ", end="", flush=True)
assistant_open = True
print(update.text, end="", flush=True)
print(f"\n\nCompleted in {iterations} iteration(s).")
async def main() -> None:
# A single credential is reused; the judge uses its own client instance.
async with AzureCliCredential() as credential:
client = FoundryChatClient(credential=credential)
judge_client = FoundryChatClient(credential=credential)
await judge_loop(client, judge_client)
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output (abridged; exact text varies by model):
=== ChatClient judge (loop until the request is answered) ===
assistant: The sky is blue because shorter (blue) wavelengths scatter more (Rayleigh scattering).
user: An evaluator reviewed your previous response and judged that it does not yet fully
address the original request.
Evaluator feedback: The response does not mention the moon.
Revise and continue so the original request is fully addressed.
assistant: The sky is blue because shorter (blue) wavelengths scatter more. At sunset, light travels
through more atmosphere, scattering away blue and leaving red/orange hues. The moon follows the
sky's colors because the same scattering applies to the light reaching it.
Completed in 2 iteration(s).
"""
@@ -1,121 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import Agent, AgentLoopMiddleware, AgentResponse
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
Agent Loop Middleware: refinement loop (should_continue + feedback tracking)
This sample demonstrates ``AgentLoopMiddleware`` driven by a ``should_continue`` predicate. The loop
keeps refining a candidate answer until the agent's latest response contains a completion marker. It
also shows feedback tracking: ``record_feedback`` logs per-iteration progress that is fed into the
next pass, ``fresh_context`` restarts each pass from the original task plus that log, and
``max_iterations`` bounds the loop as a safety cap.
``next_message`` controls the input for the next iteration (it defaults to a short "continue" nudge).
The loop is run with streaming, so the injected messages between iterations show up as ``user``
updates; the stream is printed as ``<role>: <content>`` lines.
Environment variables:
FOUNDRY_PROJECT_ENDPOINT — Azure AI Foundry project endpoint URL
FOUNDRY_MODEL — Model deployment name
Authentication:
Run ``az login`` before running this sample.
"""
COMPLETE_MARKER = "<promise>COMPLETE</promise>"
async def refinement_loop(client: FoundryChatClient) -> None:
"""Loop while the response does not yet contain a completion marker."""
print("\n=== Refinement loop (should_continue marker + feedback tracking, capped at 5) ===")
# 1. ``should_continue`` keeps the loop running until the agent signals it is done by including
# the completion marker in its latest response. It is called with the loop keyword args and
# returns True to run the agent again.
def should_continue(*, last_result: AgentResponse, **kwargs: object) -> bool:
return COMPLETE_MARKER not in last_result.text
# 2. ``record_feedback`` captures a short progress entry each iteration. Returning a string
# appends it to the log (returning None falls back to the response text). The accumulated log
# is injected into the next iteration's input so the agent builds on prior work.
def record_feedback(*, iteration: int, last_result: AgentResponse, **kwargs: object) -> str:
return f"iteration {iteration}: {last_result.text.strip()[:80]}"
# 3. ``fresh_context=True`` restarts each pass from the original task plus the progress log, and
# ``max_iterations`` bounds the loop as a safety cap.
loop = AgentLoopMiddleware(
should_continue,
max_iterations=5,
record_feedback=record_feedback,
fresh_context=True,
)
# 4. Attach the middleware to the agent.
agent = Agent(
client=client,
name="refiner",
instructions=(
"You are iteratively refining a product name for a note-taking app. Each turn, build on the "
"progress log: propose an improved candidate with a short reason. When you are confident the "
f"name is final, end your message with the exact marker {COMPLETE_MARKER}."
),
middleware=[loop],
)
# 5. Run once with streaming. The middleware drives the iterations, feeding progress forward until
# the agent emits the completion marker or the iteration cap is reached. Each contiguous
# ``user`` block marks the boundary into the next iteration, so we count loop iterations by
# those boundaries (robust to function calling, where one iteration may issue several model
# calls; tool calls/results are never ``user`` updates).
iterations = 1
in_user_block = False
assistant_open = False
async for update in agent.run("Suggest a name for a note-taking app.", stream=True):
if update.role == "user":
if not in_user_block:
iterations += 1
in_user_block = True
assistant_open = False
print(f"\nuser: {update.text}", flush=True)
continue
in_user_block = False
if update.text:
if not assistant_open:
print("\nassistant: ", end="", flush=True)
assistant_open = True
print(update.text, end="", flush=True)
print(f"\n\nCompleted in {iterations} iteration(s).")
async def main() -> None:
async with AzureCliCredential() as credential:
client = FoundryChatClient(credential=credential)
await refinement_loop(client)
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output (abridged; exact text varies by model):
=== Refinement loop (should_continue marker + feedback tracking, capped at 5) ===
assistant: "QuickJot" — short and evokes fast capture.
user: Suggest a name for a note-taking app.
user: Progress so far:
- iteration 1: "QuickJot" — short and evokes fast capture.
user: Continue working on the task. If it is complete, say so.
assistant: How about "MarginNote" — it evokes jotting ideas in the margins. <promise>COMPLETE</promise>
Completed in 2 iteration(s).
"""
@@ -1,208 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import (
Agent,
AgentLoopMiddleware,
AgentSession,
TodoProvider,
todos_remaining,
)
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
Agent Loop Middleware: todo list + report-style judge, composed as two middleware
This sample demonstrates a more complex ``AgentLoopMiddleware`` setup that composes TWO separate loop
middleware on a single agent — rather than hand-writing one predicate that does both checks. The
agent's ``middleware`` list is the composition point:
middleware=[judge_loop, todo_loop]
Agent middleware run outermost-first, so ``judge_loop`` wraps ``todo_loop``:
1. ``todo_loop`` (inner) — built from the ``todos_remaining`` helper over a ``TodoProvider``. It
re-runs the agent while any todo item is still open, so the agent plans the report and then drafts
it one todo at a time. Its final todo assembles and emits the complete report, so when the inner
loop stops its final response is the full report.
2. ``judge_loop`` (outer) — built from ``AgentLoopMiddleware.with_judge``. Each time the inner todo
loop finishes, a separate "editor" chat client reviews the assembled report (via a ``JudgeVerdict``
structured output) against a list of report ``criteria``. While the editor is not satisfied, the
outer loop re-runs the inner todo loop (the todos are already complete, so it runs the agent once)
with the editor's reasoning fed back, and the agent revises the full report.
``with_judge(criteria=...)`` renders the criteria into both the editor's judge instructions and an
extra instruction injected for the agent, so the agent writes toward the same bar the editor grades
against. A custom report-style ``instructions`` string frames the judge as an editor reviewing a
report.
The loop is run with streaming, so the injected messages between iterations show up as ``user``
updates; the stream is printed as ``<role>: <content>`` lines. Each contiguous ``user`` block (from
either loop) marks a boundary into another agent run, so the printed count is the total number of
agent runs across both loops.
Environment variables:
FOUNDRY_PROJECT_ENDPOINT — Azure AI Foundry project endpoint URL
FOUNDRY_MODEL — Model deployment name
Authentication:
Run ``az login`` before running this sample.
"""
# Requirements the finished report must satisfy. Passed as ``criteria`` to ``with_judge``, which
# renders them into both the editor's judge instructions and an extra instruction for the agent.
REPORT_REQUIREMENTS = [
"Opens with a one-paragraph executive summary.",
"Has a clearly titled section for each part of the brief.",
"Ends with a short 'Key takeaways' bulleted list.",
"Is written in clear, professional prose.",
]
# Report-style judge instructions. The ``{{criteria}}`` placeholder is replaced by ``with_judge``
# with the rendered REPORT_REQUIREMENTS block.
EDITOR_INSTRUCTIONS = (
"You are a senior editor reviewing a research report. You are given the user's original brief and "
"the report the agent produced. Decide whether the report is publication-ready. Set 'answered' to "
"true only if the report is ready, otherwise set it to false and use 'reasoning' to state "
"concisely what is missing.{{criteria}}"
)
async def report_loop(client: FoundryChatClient, editor_client: FoundryChatClient) -> None:
"""Compose a todo loop (inner) and a report-style judge loop (outer) on one agent."""
print("\n=== Todo list + report-style judge (two composed middleware) ===")
# 1. A TodoProvider gives the agent tools to plan and track the report as todo items. A single
# session (created below) keeps this todo state alive across loop iterations.
todo_provider = TodoProvider()
# 2. Inner loop: re-run the agent while the TodoProvider still has open items. ``todos_remaining``
# builds the ``should_continue`` predicate; ``max_iterations`` caps planning + one-todo-per-turn
# drafting + the final assembly turn.
todo_loop = AgentLoopMiddleware(
todos_remaining(todo_provider),
max_iterations=8,
)
# 3. Outer loop: each time the inner todo loop finishes, ``editor_client`` judges the assembled
# report against REPORT_REQUIREMENTS and the loop re-runs the inner loop while it is not yet
# publication-ready. ``with_judge`` injects the criteria for the agent too, and feeds the
# editor's reasoning back as the next iteration's input. The judge cap bounds the revision rounds.
judge_loop = AgentLoopMiddleware.with_judge(
editor_client,
instructions=EDITOR_INSTRUCTIONS,
criteria=REPORT_REQUIREMENTS,
max_iterations=4,
)
# 4. Compose the two middleware on the agent. Order matters: ``judge_loop`` is outermost (it wraps
# and re-runs the whole ``todo_loop``), ``todo_loop`` is innermost (it drives the per-todo
# drafting). The agent is told to finish with a dedicated assembly todo so that, when the inner
# loop stops, its final response is the complete report the editor then grades.
agent = Agent(
client=client,
name="report-writer",
instructions=(
"You are a research writer producing a short report. "
"On your FIRST turn, break the report into todo items using your todo tools: one item per "
"report section, plus a final 'Assemble and output the complete report' item — then stop, "
"do not start writing yet. On EACH SUBSEQUENT turn while todos remain, complete exactly "
"ONE remaining todo item, draft its content, and mark it done using your tools — never "
"more than one item per turn. When you reach the final assembly item, output the FULL "
"report in a single message and mark it done. If an editor later returns feedback, revise "
"and output the full report again."
),
context_providers=[todo_provider],
middleware=[judge_loop, todo_loop],
)
# 5. Run once with streaming. Reuse a single session so todo state persists across iterations.
# Each contiguous ``user`` block marks a boundary into another agent run; both loops inject
# such blocks (todo nudges and editor feedback), so the count is the total number of agent runs.
session = AgentSession()
prompt = "Write a brief report on the benefits and risks of remote work for software teams."
runs = 1
in_user_block = False
assistant_open = False
async for update in agent.run(prompt, session=session, stream=True):
if update.role == "user":
if not in_user_block:
runs += 1
in_user_block = True
assistant_open = False
print(f"\nuser: {update.text}", flush=True)
continue
in_user_block = False
if update.text:
if not assistant_open:
print("\nassistant: ", end="", flush=True)
assistant_open = True
print(update.text, end="", flush=True)
print(f"\n\nCompleted in {runs} agent run(s).")
# 6. Inspect the todos the agent created, loaded from the same store the inner loop uses.
items = await todo_provider.store.load_items(session, source_id=todo_provider.source_id)
print("\nTodos after the run:")
for item in items:
mark = "x" if item.is_complete else " "
print(f" [{mark}] {item.id}. {item.title}")
"""
Sample output for ``report_loop`` (abridged; exact text varies by model):
=== Todo list + report-style judge (two composed middleware) ===
assistant: Here is my plan. I'll create todos for each section and a final assembly item.
user: Continue working on the task. If it is complete, say so.
...
assistant: # Remote Work for Software Teams
**Executive summary:** Remote work offers flexibility and access to wider talent...
## Benefits
...
## Risks
...
## Key takeaways
- Flexibility improves retention.
- Async communication needs discipline.
user: An evaluator reviewed your previous response and judged that it does not yet fully
address the original request.
Evaluator feedback: Add a one-paragraph executive summary before the first section.
Revise and continue so the original request is fully addressed.
assistant: # Remote Work for Software Teams
**Executive summary:** ... (revised, now opens with a summary)
...
Completed in 7 agent run(s).
Todos after the run:
[x] 1. Benefits section
[x] 2. Risks section
[x] 3. Key takeaways
[x] 4. Assemble and output the complete report
"""
async def main() -> None:
# A single credential is reused; the editor judge uses its own client instance.
async with AzureCliCredential() as credential:
client = FoundryChatClient(credential=credential)
editor_client = FoundryChatClient(credential=credential)
await report_loop(client, editor_client)
if __name__ == "__main__":
asyncio.run(main())
@@ -1,129 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
from agent_framework import Agent, AgentLoopMiddleware, AgentSession, TodoProvider, todos_remaining
from agent_framework.foundry import FoundryChatClient
from azure.identity.aio import AzureCliCredential
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()
"""
Agent Loop Middleware: todo loop (should_continue via a provider helper)
This sample demonstrates ``AgentLoopMiddleware`` driven by a ``should_continue`` predicate built from
a ``TodoProvider``. The ``todos_remaining`` helper keeps the agent running while it still has open
todo items, so the agent plans work on its first turn and completes one item per turn afterwards.
``max_iterations`` bounds the loop as a safety cap, and a single session keeps the todo state across
iterations. After the run the sample prints the todos the agent created.
The loop is run with streaming, so the injected messages between iterations show up as ``user``
updates; the stream is printed as ``<role>: <content>`` lines.
Environment variables:
FOUNDRY_PROJECT_ENDPOINT — Azure AI Foundry project endpoint URL
FOUNDRY_MODEL — Model deployment name
Authentication:
Run ``az login`` before running this sample.
"""
async def todo_loop(client: FoundryChatClient) -> None:
"""Loop while a TodoProvider still has open items."""
print("\n=== Callable criterion (loop while todos remain) ===")
# 1. A TodoProvider gives the agent tools to plan and track work as todo items.
todo_provider = TodoProvider()
# 2. ``todos_remaining`` builds a ``should_continue`` predicate that returns True while any todo
# item is still open. ``max_iterations`` guarantees the loop stops even if the agent stalls.
loop = AgentLoopMiddleware(
should_continue=todos_remaining(todo_provider),
max_iterations=6,
)
agent = Agent(
client=client,
name="planner",
instructions=(
"You are a writing assistant working through a todo list. "
"On your FIRST turn, break the task into todo items using your todo tools and stop "
"(do not start writing yet). On EACH SUBSEQUENT turn, complete exactly ONE remaining "
"todo item, write its content, and mark it done using your tools — never complete more "
"than one item per turn. When every item is done, give a brief final summary."
),
context_providers=[todo_provider],
middleware=[loop],
)
# 3. Reuse a single session so todo state persists across loop iterations. Each contiguous
# ``user`` block marks the boundary into the next iteration, so we count loop iterations by
# those boundaries — robust to the function calling this loop relies on (the todo tools issue
# several model calls per iteration, but tool calls/results are never ``user`` updates).
session = AgentSession()
prompt = "Plan and write a short 3-section blog post about Rayleigh scattering."
iterations = 1
in_user_block = False
assistant_open = False
async for update in agent.run(prompt, session=session, stream=True):
if update.role == "user":
if not in_user_block:
iterations += 1
in_user_block = True
assistant_open = False
print(f"\nuser: {update.text}", flush=True)
continue
in_user_block = False
if update.text:
if not assistant_open:
print("\nassistant: ", end="", flush=True)
assistant_open = True
print(update.text, end="", flush=True)
print(f"\n\nCompleted in {iterations} iteration(s).")
# 4. Inspect the todos the agent created, loaded from the same store the loop predicate uses.
items = await todo_provider.store.load_items(session, source_id=todo_provider.source_id)
print("\nTodos after the run:")
for item in items:
mark = "x" if item.is_complete else " "
print(f" [{mark}] {item.id}. {item.title}")
async def main() -> None:
async with AzureCliCredential() as credential:
client = FoundryChatClient(credential=credential)
await todo_loop(client)
if __name__ == "__main__":
asyncio.run(main())
"""
Sample output (abridged; exact text varies by model):
=== Callable criterion (loop while todos remain) ===
assistant: Here is my plan. I'll create todos for each section.
user: Progress so far:
- Here is my plan. I'll create todos for each section.
user: Continue working on the task. If it is complete, say so.
assistant: Section 1 drafted. Marking it done.
user: Progress so far:
- Section 1 drafted. Marking it done.
user: Continue working on the task. If it is complete, say so.
assistant: Section 2 drafted. Marking it done.
user: Progress so far:
- Section 2 drafted. Marking it done.
user: Continue working on the task. If it is complete, say so.
assistant: Section 3 drafted. Marking it done.
Completed in 4 iteration(s).
Todos after the run:
[x] 1. Draft "What light is" section
[x] 2. Draft "How Rayleigh scattering works" section
[x] 3. Draft "Why the sky is blue" section
"""
@@ -40,12 +40,12 @@ async def main() -> None:
print(f"User: {message}")
if stream:
print("Assistant: ", end="")
async for chunk in client.get_response(messages, options={"tools": [get_time]}, stream=True):
async for chunk in client.get_response(messages, tools=get_time, stream=True):
if str(chunk):
print(str(chunk), end="")
print("")
else:
response = await client.get_response(messages, options={"tools": [get_time]})
response = await client.get_response(messages, tools=get_time)
print(f"Assistant: {response}")
@@ -70,7 +70,9 @@ async def run_policy_flow(
("good (warm cache)", GOOD_PROMPT_FOLLOWUP),
]
for tag, text in prompts:
response: AgentResponse = await agent.run(Message("user", [text], additional_properties={"user_id": user_id}))
response: AgentResponse = await agent.run(
Message("user", [text], additional_properties={"user_id": user_id})
)
outcome = "BLOCKED" if blocked_marker in str(response).lower() else "ALLOWED"
print(f"[{label}] {tag}: {outcome}\n{response}\n")
@@ -205,7 +207,9 @@ async def run_with_chat_middleware() -> None:
model=deployment,
project_endpoint=endpoint,
credential=AzureCliCredential(),
middleware=[PurviewChatPolicyMiddleware(build_credential(), settings)],
middleware=[
PurviewChatPolicyMiddleware(build_credential(), settings)
],
)
agent = Agent(