Compare commits

...
Author SHA1 Message Date
dependabot[bot]andGitHub 6b2ff3ed7a Build(deps): Bump python-multipart from 0.0.26 to 0.0.27 in /python
Bumps [python-multipart](https://github.com/Kludex/python-multipart) from 0.0.26 to 0.0.27.
- [Release notes](https://github.com/Kludex/python-multipart/releases)
- [Changelog](https://github.com/Kludex/python-multipart/blob/main/CHANGELOG.md)
- [Commits](https://github.com/Kludex/python-multipart/compare/0.0.26...0.0.27)

---
updated-dependencies:
- dependency-name: python-multipart
  dependency-version: 0.0.27
  dependency-type: indirect
...

Signed-off-by: dependabot[bot] <support@github.com>
2026-06-09 06:01:18 +00:00
bad05a2bdc Python: Harness console for python (#6312)
* Add initial harness console for python

* Add textual to project

* Add planning and approval flows with list selector

* Address PR comments

* Fix list selection bug

* Fix PR #6312 round 2 review comments

- Escape untrusted agent text with rich.markup.escape() in observers
  (text_output, planning_output, reasoning_display) to prevent markup injection
- Remove non-functional 'Always approve' choices from tool_approval.py
  (framework lacks CreateAlwaysApproveToolResponse support)
- Remove textual from root pyproject.toml dev deps (sample-specific)
- Add PEP 723 inline script metadata to harness_research.py
- Narrow except Exception to except NoMatches in list_selection.py

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix build error

* Fix build errors

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-09 05:48:35 +00:00
7e0767a0a0 Python: Fix per-service-call history persistence with server-storing clients (#6310)
* Fix per-service-call history persistence with server-storing clients

When an Agent set require_per_service_call_history_persistence=True together
with a HistoryProvider, and the chat client stored history server-side by
default (e.g. OpenAIChatClient, STORES_BY_DEFAULT=True), the external history
provider was silently never persisted.

Unify persistence on the per-service-call middleware: when the flag is set and
a HistoryProvider exists, the middleware is always installed and owns
persistence. service_stores_history now only selects middleware behavior:
- service does not store: load providers and drive the function loop with a
  local sentinel conversation id, or
- service stores: skip loading (the service owns history) and persist each
  service call while the real conversation id flows through.

Also rationalize chat-options handling in _prepare_run_context:
- _merge_options now skips None overrides and strips remaining None values, so
  an unset `store` is never forwarded and the service decides its own default.
- Resolve `store` and `conversation_id` once from a single combined view
  (effective_options) instead of probing both default and runtime dicts; the
  auto-injection and per-service-call resolution now agree on conversation_id.

Fixes #5798

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Correct as_agent() docstring: persistence is per service call, not once per run

Address PR review: when the client stores history server-side, the
per-service-call middleware still persists after each model call; only
provider loading is skipped. The previous "persist once per run()" wording
contradicted the implementation.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Address PR review: docs, missing-conversation-id warning, and tests

- Clarify that require_per_service_call_history_persistence is a no-op when no
  HistoryProvider is present (docstrings in _agents.py and _clients.py).
- Warn on every service call when the client stores history server-side but
  returns no conversation_id, so the (uncommon) loss of cross-turn resumability
  cannot fail silently.
- Add tests: storing client + existing conversation_id does not raise and the id
  propagates; two runs on the same session keep persisting with a stable
  service_session_id and no provider loading; storing-without-conversation-id
  warns per call.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-09 05:47:57 +00:00
af772997af .NET: [BREAKING] Migrate .NET GitHub Copilot SDK to v1.0.0 (#6381)
* Migrate .NET GitHub Copilot SDK from 1.0.0-beta.2 to 1.0.0

- Update namespace from GitHub.Copilot.SDK to GitHub.Copilot
- Replace PermissionRequestResult/PermissionRequestResultKind with PermissionDecision
- Remove ConnectionState check (StartAsync is now idempotent)
- Rename ConfigDir to ConfigDirectory
- Use SessionConfig.Clone() for CopySessionConfig
- Update Tools type from List<AIFunction> to List<AIFunctionDeclaration>
- Rename UserMessageAttachmentFile to AttachmentFile
- Update usage data types (CacheWriteTokens: long, Duration: TimeSpan)
- Add GHCP001 NoWarn for experimental SDK APIs (matches framework convention)
- Specify type argument on CopilotSession.On<SessionEvent>()

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix formatting: remove unused using directive

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Skip AzureFunctions SamplesValidation tests pending func tools fix

Azure Functions Core Tools v4 can no longer auto-detect the worker
runtime in CI (local.settings.json is gitignored). All 7 active
SamplesValidation tests fail with 'Worker runtime cannot be None'.

Tracked by: https://github.com/microsoft/agent-framework/issues/6402

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Skip additional failing integration tests in CI

WorkflowSamplesValidation (5 tests): same func tools issue as #6402.
WorkflowConsoleAppSamplesValidation (4 tests): KeyNotFoundException
during workflow execution, tracked by #6404.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-08 22:34:05 +00:00
westeyandGitHub b343625c1f .NET: Add approval bypassing to harness as the default (#6387)
* Add approval bypassing to harness as a default

* Add tests

* Address PR comments.
2026-06-08 17:50:41 +00:00
Evan MattsonandGitHub 9bc7b27813 Match AG-UI approval responses to requested arguments (#6376) 2026-06-08 16:33:16 +00:00
westeyandGitHub 6a2efeae7c .NET: [BREAKING] Fix hosting bugs (#6388)
* Fix hosting bugs

* Address PR comments
2026-06-08 16:17:54 +00:00
Vedant SonaniandGitHub 6169df04cb Python: fix(mem0): isolate entity retrieval and correct app_id payload (#6242)
* fix(mem0): parallel memory retrieval logic and strict type compliance

* fix(mem0): align parallel retrieval types for pyright and mypy

* fix(mem0): handle asyncio.CancelledError in search response and update test description

* fix(mem0): improve error handling for asyncio.CancelledError and update test names for clarity

* fix(mem0): improve retrieval response handling
2026-06-08 13:50:23 +00:00
Peter IbekweandGitHub 331201294b .NET: Fix single-column value unwrap in declarative workflow (#6367)
* Fix single-column value unwrap in declarative workflow

* Added more tests
2026-06-08 11:37:12 +00:00
Yufeng HeandGitHub fa9e086576 fix: preserve foreach record values (#6208) 2026-06-05 22:01:59 +00:00
dcc218dbac Python: feat(python): Add MCP client OTel spans per GenAI semantic conventions (#6349)
* feat(python): Add MCP client OTel spans per GenAI semantic conventions

Implement MCP client spans per the OTel GenAI Semantic Conventions for MCP
(https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/#client).

Operations instrumented:
- initialize: CLIENT span capturing MCP session setup
- tools/list: CLIENT span for tool listing (per-page)
- prompts/list: CLIENT span for prompt listing (per-page)
- tools/call: CLIENT span (nested under execute_tool when called via FunctionTool)
- prompts/get: CLIENT span

Span attributes follow the MCP semantic conventions:
- Required: mcp.method.name
- Conditional: error.type, gen_ai.tool.name, gen_ai.prompt.name
- Recommended: gen_ai.operation.name, mcp.protocol.version, mcp.session.id,
  network.transport, server.address, server.port

Transport-specific attributes per subclass:
- MCPStdioTool: network.transport=pipe
- MCPStreamableHTTPTool: network.transport=tcp, network.protocol.name=http
- MCPWebsocketTool: network.transport=tcp, network.protocol.name=websocket

All span creation gated behind OBSERVABILITY_SETTINGS.ENABLED.

Closes #3624
Closes #4697

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* refactor: simplify MCP spans — remove enrichment logic and protocol version caching

- Always create nested CLIENT spans for tools/call instead of enriching
  the parent execute_tool span
- Remove _ACTIVE_TOOL_EXECUTION_SPAN contextvar (no longer needed)
- Remove enrich_span_with_mcp_attributes() helper
- Remove _otel_error_type preservation in FunctionTool.invoke()
- Remove _mcp_protocol_version instance variable; protocol version is
  only set on the initialize span where it is available

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Refine copilot solution

* fix: enable automatic exception recording on MCP spans

Remove record_exception=False and set_status_on_exception=False from
create_mcp_client_span. Let OTel handle exception recording and status
setting automatically. The manual set_mcp_span_error calls for tools/call
still correctly set error.type (which OTel's automatic handling doesn't
touch), so tool_error is preserved.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Reduce number of lines

* Add comment to sample

* test: address PR review comments on MCP observability tests

- Fix initialize test to call mocked session.initialize() and read
  protocolVersion from the result instead of hardcoding it
- Add tools/call McpError error-path test
- Add prompts/get McpError error-path test

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix export error

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-05 19:23:01 +00:00
6bd2cfec03 .NET: [BREAKING] Add auto-approval rules (heuristics) to ToolApprovalAgent (#6335)
* Add support for approving tools via heuristic rules

* Address PR comments

* Address PR comments

* Apply suggestion from @SergeyMenshykh

Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>

---------

Co-authored-by: SergeyMenshykh <68852919+SergeyMenshykh@users.noreply.github.com>
2026-06-05 18:43:07 +01:00
westeyandGitHub ab8ba8fc61 .NET: Allow storage of auto-approved functions (#4950)
* Allow storage of auto-approved functions

* Address PR comments
2026-06-05 18:42:21 +01:00
Tao ChenandGitHub 9cafd7e58b Python: Refactor workflow as agent pending request handling (#6259)
* WIP: Refactor Workflow as agent pending request handling

* WIP: debugging empty message bug

* Working: Workflow as agent with function approval

* Address Copilot comments

* Fix mypy

* Address comments and fix pipeline

* Request info non function approval now becomes function call

* Revert uv.lock

* Fix mypy

* Bump min version of azure-ai-project

* Remove RequestInfoFunctionArgs

* fix tests

* Fix failing tests

* Fix sample
2026-06-05 17:23:19 +00:00
d5335fbeae Python (fix:gemini): make Gemini honor declarative outputSchema, not just JSON mode (#5893)
* fix(gemini): preserve schema response_format

* fix(gemini): satisfy pyright strict in response schema extraction

Cast Any-narrowed mappings to Mapping[str, Any] in the structured-output
schema helpers so pyright strict no longer reports partially-unknown
member, argument, and variable types. Pass response_format["format"]
straight into the recursive extractor, which already guards non-mapping
inputs. No behavior change.

* fix(gemini): use Sequence[object] cast to satisfy both mypy and pyright

The Sequence[Any] cast pyright strict needs to know the loop element type
is reported as a redundant-cast by mypy, which already narrows the
isinstance branch to Sequence[Any]. Cast to Sequence[object] instead:
pyright gets a fully known element type and mypy no longer sees an
identical-type cast. No behavior change.

---------

Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>
2026-06-05 15:17:51 +00:00
bf4ad48cf2 Python: MCP long-running task support in Python (#6319)
* MCP long-running task support in Python

* Fix pyupgrade and AGENTS.md reconnect description

- pyupgrade: drop forward-reference string annotations in _mcp.py (Python 3.10+ resolves them natively now that MCPTaskOptions is defined before use).

- AGENTS.md: align reconnect description with current behavior. Phase 1 (initial tools/call) does NOT retry on connection loss; raises 'connection lost; task state unknown' instead, so a server that accepted the request but lost the response cannot start the operation twice. Phase 2 (tasks/get / tasks/result) still reconnects once against the same task_id.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* Fix bandit nosec marker for CI pipeline

* Address PR feedbacks

* Clarifiied comments and addressed more PR feedbacks.

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
2026-06-05 00:04:55 +00:00
112 changed files with 12546 additions and 780 deletions
+1 -1
View File
@@ -99,7 +99,7 @@
<PackageVersion Include="Microsoft.SemanticKernel.Connectors.InMemory" Version="1.67.0-preview" />
<PackageVersion Include="Microsoft.SemanticKernel.Connectors.Qdrant" Version="1.67.0-preview" />
<!-- Agent SDKs -->
<PackageVersion Include="GitHub.Copilot.SDK" Version="1.0.0-beta.2" />
<PackageVersion Include="GitHub.Copilot.SDK" Version="1.0.0" />
<PackageVersion Include="Microsoft.Agents.CopilotStudio.Client" Version="1.3.171-beta" />
<!-- M365 Agents SDK -->
<PackageVersion Include="AdaptiveCards" Version="3.1.0" />
@@ -6,6 +6,7 @@
<Nullable>enable</Nullable>
<ImplicitUsings>enable</ImplicitUsings>
<NoWarn>$(NoWarn);GHCP001</NoWarn>
</PropertyGroup>
<ItemGroup>
@@ -2,21 +2,22 @@
// This sample shows how to create a GitHub Copilot agent with shell command permissions.
using GitHub.Copilot.SDK;
using GitHub.Copilot;
using GitHub.Copilot.Rpc;
using Microsoft.Agents.AI;
// Permission handler that prompts the user for approval
static Task<PermissionRequestResult> PromptPermission(PermissionRequest request, PermissionInvocation invocation)
static Task<PermissionDecision> PromptPermission(PermissionRequest request, PermissionInvocation invocation)
{
Console.WriteLine($"\n[Permission Request: {request.Kind}]");
Console.Write("Approve? (y/n): ");
string? input = Console.ReadLine()?.Trim().ToUpperInvariant();
PermissionRequestResultKind kind = input is "Y" or "YES"
? PermissionRequestResultKind.Approved
: PermissionRequestResultKind.Rejected;
PermissionDecision decision = input is "Y" or "YES"
? PermissionDecision.ApproveOnce()
: PermissionDecision.Reject();
return Task.FromResult(new PermissionRequestResult { Kind = kind });
return Task.FromResult(decision);
}
// Create and start a Copilot client
@@ -36,7 +36,7 @@ dotnet run
You can customize the agent by providing additional configuration:
```csharp
using GitHub.Copilot.SDK;
using GitHub.Copilot;
using Microsoft.Agents.AI;
// Create and start a Copilot client
@@ -44,18 +44,33 @@ public static class HostedFoundryMemoryProviderScopes
session => new FoundryMemoryProvider.State(new FoundryMemoryProviderScope(GetRequiredHostedContext(session).ChatId));
/// <summary>
/// Returns a <c>stateInitializer</c> that scopes memories per (user, chat) pair, using
/// <c>"{UserId}:{ChatId}"</c> as the partition key. Use this when memories should be visible
/// only to the same user within the same conversation.
/// Returns a <c>stateInitializer</c> that scopes memories per (user, chat) pair, composing
/// <see cref="HostedSessionContext.UserId"/> and <see cref="HostedSessionContext.ChatId"/> into a
/// single delimiter-safe partition key. Use this when memories should be visible only to the same
/// user within the same conversation.
/// </summary>
/// <remarks>
/// Both identity values are opaque strings that may contain any characters, including the <c>:</c>
/// delimiter. To keep the composite key injective (so two distinct (user, chat) pairs can never
/// collide), each part is escaped (<c>\</c> becomes <c>\\</c>, then <c>:</c> becomes <c>\:</c>) before
/// being joined with a <c>::</c> separator.
/// </remarks>
/// <returns>A delegate suitable for the <c>stateInitializer</c> argument of <see cref="FoundryMemoryProvider"/>.</returns>
public static Func<AgentSession?, FoundryMemoryProvider.State> PerUserAndChat() =>
session =>
{
var ctx = GetRequiredHostedContext(session);
return new FoundryMemoryProvider.State(new FoundryMemoryProviderScope($"{ctx.UserId}:{ctx.ChatId}"));
return new FoundryMemoryProvider.State(
new FoundryMemoryProviderScope($"{EscapeScopePart(ctx.UserId)}::{EscapeScopePart(ctx.ChatId)}"));
};
/// <summary>
/// Escapes special characters in a scope part so that distinct (user, chat) pairs produce distinct
/// composite scope keys. Backslashes are escaped first (<c>\</c> becomes <c>\\</c>), then colons
/// (<c>:</c> becomes <c>\:</c>), ensuring the <c>{user}::{chat}</c> format is unambiguous.
/// </summary>
private static string EscapeScopePart(string part) => part.Replace("\\", "\\\\").Replace(":", "\\:");
private static HostedSessionContext GetRequiredHostedContext(AgentSession? session) =>
session?.GetHostedContext()
?? throw new InvalidOperationException(
@@ -6,7 +6,7 @@ using Microsoft.Agents.AI.GitHub.Copilot;
using Microsoft.Extensions.AI;
using Microsoft.Shared.Diagnostics;
namespace GitHub.Copilot.SDK;
namespace GitHub.Copilot;
/// <summary>
/// Provides extension methods for <see cref="CopilotClient"/>
@@ -9,7 +9,7 @@ using System.Text.Json;
using System.Threading;
using System.Threading.Channels;
using System.Threading.Tasks;
using GitHub.Copilot.SDK;
using GitHub.Copilot;
using Microsoft.Extensions.AI;
using Microsoft.Shared.Diagnostics;
@@ -169,7 +169,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
Channel<AgentResponseUpdate> channel = Channel.CreateUnbounded<AgentResponseUpdate>();
// Subscribe to session events
using IDisposable subscription = copilotSession.On(evt =>
using IDisposable subscription = copilotSession.On<SessionEvent>(evt =>
{
switch (evt)
{
@@ -210,7 +210,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
string prompt = string.Join("\n", messages.Select(m => m.Text));
// Handle DataContent as attachments
(List<UserMessageAttachmentFile>? attachments, tempDir) = await ProcessDataContentAttachmentsAsync(
(List<AttachmentFile>? attachments, tempDir) = await ProcessDataContentAttachmentsAsync(
messages,
cancellationToken).ConfigureAwait(false);
@@ -262,10 +262,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
private async Task EnsureClientStartedAsync(CancellationToken cancellationToken)
{
if (this._copilotClient.State != ConnectionState.Connected)
{
await this._copilotClient.StartAsync(cancellationToken).ConfigureAwait(false);
}
await this._copilotClient.StartAsync(cancellationToken).ConfigureAwait(false);
}
private ResumeSessionConfig CreateResumeConfig()
@@ -275,36 +272,18 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
/// <summary>
/// Copies all supported properties from a source <see cref="SessionConfig"/> into a new instance
/// with <see cref="SessionConfig.Streaming"/> set to <c>true</c>.
/// with <see cref="SessionConfigBase.Streaming"/> set to <c>true</c>.
/// </summary>
internal static SessionConfig CopySessionConfig(SessionConfig source)
{
return new SessionConfig
{
Model = source.Model,
ReasoningEffort = source.ReasoningEffort,
Tools = source.Tools,
SystemMessage = source.SystemMessage,
AvailableTools = source.AvailableTools,
ExcludedTools = source.ExcludedTools,
Provider = source.Provider,
OnPermissionRequest = source.OnPermissionRequest,
OnUserInputRequest = source.OnUserInputRequest,
Hooks = source.Hooks,
WorkingDirectory = source.WorkingDirectory,
ConfigDir = source.ConfigDir,
McpServers = source.McpServers,
CustomAgents = source.CustomAgents,
SkillDirectories = source.SkillDirectories,
DisabledSkills = source.DisabledSkills,
InfiniteSessions = source.InfiniteSessions,
Streaming = true
};
SessionConfig copy = source.Clone();
copy.Streaming = true;
return copy;
}
/// <summary>
/// Copies all supported properties from a source <see cref="SessionConfig"/> into a new
/// <see cref="ResumeSessionConfig"/> with <see cref="ResumeSessionConfig.Streaming"/> set to <c>true</c>.
/// <see cref="ResumeSessionConfig"/> with <see cref="SessionConfigBase.Streaming"/> set to <c>true</c>.
/// </summary>
internal static ResumeSessionConfig CopyResumeSessionConfig(SessionConfig? source)
{
@@ -321,7 +300,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
OnUserInputRequest = source?.OnUserInputRequest,
Hooks = source?.Hooks,
WorkingDirectory = source?.WorkingDirectory,
ConfigDir = source?.ConfigDir,
ConfigDirectory = source?.ConfigDirectory,
McpServers = source?.McpServers,
CustomAgents = source?.CustomAgents,
SkillDirectories = source?.SkillDirectories,
@@ -394,10 +373,10 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
AdditionalPropertiesDictionary<long>? additionalCounts = null;
if (usageEvent.Data.CacheWriteTokens is double cacheWriteTokens)
if (usageEvent.Data.CacheWriteTokens is long cacheWriteTokens)
{
additionalCounts ??= [];
additionalCounts[nameof(AssistantUsageData.CacheWriteTokens)] = (long)cacheWriteTokens;
additionalCounts[nameof(AssistantUsageData.CacheWriteTokens)] = cacheWriteTokens;
}
if (usageEvent.Data.Cost is double cost)
@@ -406,10 +385,10 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
additionalCounts[nameof(AssistantUsageData.Cost)] = (long)cost;
}
if (usageEvent.Data.Duration is double duration)
if (usageEvent.Data.Duration is TimeSpan duration)
{
additionalCounts ??= [];
additionalCounts[nameof(AssistantUsageData.Duration)] = (long)duration;
additionalCounts[nameof(AssistantUsageData.Duration)] = (long)duration.TotalMilliseconds;
}
return additionalCounts;
@@ -432,7 +411,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
private static SessionConfig? GetSessionConfig(IList<AITool>? tools, string? instructions)
{
List<AIFunction>? mappedTools = tools is { Count: > 0 } ? tools.OfType<AIFunction>().ToList() : null;
List<AIFunctionDeclaration>? mappedTools = tools is { Count: > 0 } ? tools.OfType<AIFunctionDeclaration>().ToList() : null;
SystemMessageConfig? systemMessage = instructions is not null ? new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = instructions } : null;
if (mappedTools is null && systemMessage is null)
@@ -443,11 +422,11 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
return new SessionConfig { Tools = mappedTools, SystemMessage = systemMessage };
}
private static async Task<(List<UserMessageAttachmentFile>? Attachments, string? TempDir)> ProcessDataContentAttachmentsAsync(
private static async Task<(List<AttachmentFile>? Attachments, string? TempDir)> ProcessDataContentAttachmentsAsync(
IEnumerable<ChatMessage> messages,
CancellationToken cancellationToken)
{
List<UserMessageAttachmentFile>? attachments = null;
List<AttachmentFile>? attachments = null;
string? tempDir = null;
foreach (ChatMessage message in messages)
{
@@ -461,7 +440,7 @@ public sealed class GitHubCopilotAgent : AIAgent, IAsyncDisposable
string tempFilePath = await dataContent.SaveToAsync(tempDir, cancellationToken).ConfigureAwait(false);
attachments ??= [];
attachments.Add(new UserMessageAttachmentFile
attachments.Add(new AttachmentFile
{
Path = tempFilePath,
DisplayName = Path.GetFileName(tempFilePath)
@@ -4,6 +4,7 @@
<VersionSuffix>preview</VersionSuffix>
<!-- GitHub.Copilot.SDK only supports .NET 8.0+ -->
<TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
<NoWarn>$(NoWarn);GHCP001</NoWarn>
</PropertyGroup>
<PropertyGroup>
@@ -138,7 +138,7 @@ public sealed class HarnessAgent : DelegatingAIAgent
if (options?.DisableToolApproval is not true)
{
builder.UseToolApproval();
builder.UseToolApproval(options?.ToolApprovalAgentOptions);
}
if (options?.DisableOpenTelemetry is not true)
@@ -178,8 +178,14 @@ public sealed class HarnessAgent : DelegatingAIAgent
IEnumerable<AIContextProvider> contextProviders = BuildContextProviders(options, loggerFactory);
return chatClient
.AsBuilder()
ChatClientBuilder chatClientBuilder = chatClient.AsBuilder();
if (options?.DisableNonApprovalRequiredFunctionBypassing is not true)
{
chatClientBuilder.UseNonApprovalRequiredFunctionBypassing();
}
return chatClientBuilder
.UseFunctionInvocation(loggerFactory, configure: options?.MaximumIterationsPerRequest is int maxIterations
? ficc => ficc.MaximumIterationsPerRequest = maxIterations
: null)
@@ -101,6 +101,29 @@ public sealed class HarnessAgentOptions
/// </remarks>
public bool DisableToolApproval { get; set; }
/// <summary>
/// Gets or sets the options for the <see cref="ToolApprovalAgent"/> middleware.
/// </summary>
/// <remarks>
/// When <see langword="null"/>, the <see cref="ToolApprovalAgent"/> uses default settings.
/// This property has no effect when <see cref="DisableToolApproval"/> is <see langword="true"/>.
/// </remarks>
public ToolApprovalAgentOptions? ToolApprovalAgentOptions { get; set; }
/// <summary>
/// Gets or sets a value indicating whether bypassing of approval requests for tools that do not
/// require approval is disabled.
/// </summary>
/// <remarks>
/// When <see langword="false"/> (the default), the underlying chat client pipeline includes the decorator
/// added by <see cref="ChatClientBuilderExtensions.UseNonApprovalRequiredFunctionBypassing"/> above the
/// function invocation middleware.
/// This stores automatically approved function calls for tools that do not require approval in the session
/// state when they are returned alongside tools that do, so that only tools that truly require human
/// approval are surfaced to the caller.
/// </remarks>
public bool DisableNonApprovalRequiredFunctionBypassing { get; set; }
/// <summary>
/// Gets or sets a value indicating whether the <see cref="FileMemoryProvider"/> is disabled.
/// </summary>
@@ -21,6 +21,18 @@ namespace Microsoft.Agents.AI.Hosting;
/// from the ambient <see cref="HttpContext"/>.
/// </para>
/// <para>
/// <strong>Security warning:</strong> The configured <see cref="ClaimsIdentitySessionIsolationKeyProviderOptions.ClaimType"/>
/// must uniquely identify the principal within the served population. Display names, usernames, email
/// aliases, and other mutable or non-unique claims are <strong>unsafe</strong> isolation keys unless the
/// host can prove their uniqueness across all callers: two distinct principals that share the same value
/// would receive the same isolation key and could read or overwrite one another's persisted sessions.
/// The default claim type is <see cref="ClaimTypes.NameIdentifier"/>, a stable unique subject identifier
/// that is typically populated from the OpenID Connect <c>sub</c> claim via the default JWT inbound claim
/// mapping (note that this differs from Entra's object identifier <c>oid</c> claim; override
/// <see cref="ClaimsIdentitySessionIsolationKeyProviderOptions.ClaimType"/> if you need <c>oid</c> or your
/// provider maps a different claim).
/// </para>
/// <para>
/// If the <see cref="HttpContext"/> is unavailable, the user is not authenticated, or the specified claim
/// is missing, the provider returns <see langword="null"/>. The consuming <see cref="IsolationKeyScopedAgentSessionStore"/>
/// will then enforce strict or pass-through behavior based on its configuration.
@@ -60,18 +72,24 @@ public class ClaimsIdentitySessionIsolationKeyProvider : SessionIsolationKeyProv
/// <param name="cancellationToken">The <see cref="CancellationToken"/> to monitor for cancellation requests.</param>
/// <returns>
/// A task that represents the asynchronous operation. The task result contains the value of the
/// configured claim type from the current user's identity, or <see langword="null"/> if the claim
/// is not present or the HTTP context is unavailable.
/// configured claim type from the current user's identity, or <see langword="null"/> if the HTTP
/// context is unavailable, the user is not authenticated, or the claim is not present.
/// </returns>
/// <remarks>
/// This method retrieves the claim value from <c>HttpContext.User.Claims</c>. If multiple claims
/// of the specified type exist, the first match is returned.
/// This method only reads claims from an authenticated principal: if the current request has no
/// authenticated user, it returns <see langword="null"/> rather than trusting claims on an
/// unauthenticated identity. The claim value is retrieved from <c>HttpContext.User.Claims</c>; if
/// multiple claims of the specified type exist, the first match is returned.
/// </remarks>
public override ValueTask<string?> GetSessionIsolationKeyAsync(CancellationToken cancellationToken = default)
{
Claim? claim = this._httpContextAccessor?
.HttpContext?
.User?.Claims.FirstOrDefault(c => c.Type == this._claimType);
ClaimsPrincipal? user = this._httpContextAccessor?.HttpContext?.User;
if (user?.Identity?.IsAuthenticated != true)
{
return new ValueTask<string?>((string?)null);
}
Claim? claim = user?.Claims.FirstOrDefault(c => c.Type == this._claimType);
return new ValueTask<string?>(claim?.Value);
}
@@ -14,17 +14,30 @@ public class ClaimsIdentitySessionIsolationKeyProviderOptions
/// </summary>
/// <remarks>
/// <para>
/// Defaults to <see cref="ClaimsIdentity.DefaultNameClaimType"/>, which typically corresponds to
/// the user's name or unique identifier claim.
/// Defaults to <see cref="ClaimTypes.NameIdentifier"/>, which corresponds to a stable, unique
/// subject identifier for the authenticated principal. For OpenID Connect tokens (including those
/// issued by Microsoft Entra ID), this is typically populated from the <c>sub</c> claim via the
/// default JWT inbound claim mapping. Note that <c>sub</c> is distinct from Entra's object
/// identifier (<c>oid</c>) claim; if you require the <c>oid</c> claim, or your provider does not map
/// a unique identifier onto <see cref="ClaimTypes.NameIdentifier"/>, override <see cref="ClaimType"/>
/// with the appropriate claim type.
/// </para>
/// <para>
/// <strong>Security warning:</strong> The configured claim must uniquely identify the principal
/// within the served population. Display names (<see cref="ClaimsIdentity.DefaultNameClaimType"/>
/// / <see cref="ClaimTypes.Name"/>), usernames, email aliases, and other mutable or non-unique
/// claims are <strong>unsafe</strong> isolation keys unless the host can prove their uniqueness
/// across all callers. Two distinct principals that share the same value for a non-unique claim
/// would receive the same session-isolation key and could read or overwrite one another's
/// persisted sessions. Only override this value with a claim that is guaranteed unique and stable.
/// </para>
/// <para>
/// Common alternatives include:
/// <list type="bullet">
/// <item><description><c>ClaimTypes.NameIdentifier</c> — Stable user identifier</description></item>
/// <item><description><c>ClaimTypes.Email</c> — Email address</description></item>
/// <item><description>Custom claim types specific to your authentication provider</description></item>
/// <item><description>A composite of tenant and subject identifiers — required for multi-tenant hosts where the subject is only unique per tenant</description></item>
/// <item><description>Custom claim types specific to your authentication provider, provided they are unique and stable</description></item>
/// </list>
/// </para>
/// </remarks>
public string ClaimType { get; set; } = ClaimsIdentity.DefaultNameClaimType;
public string ClaimType { get; set; } = ClaimTypes.NameIdentifier;
}
@@ -1,6 +1,7 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Security.Claims;
using Microsoft.AspNetCore.Http;
using Microsoft.Extensions.DependencyInjection;
@@ -19,8 +20,28 @@ public static class ServiceCollectionExtensions
/// <param name="options"> Optional configuration for the claims-based session isolation key provider.</param>
/// <returns>The <see cref="IServiceCollection"/> so that additional calls can be chained.</returns>
/// <remarks>
/// <para>
/// This method requires <see cref="IHttpContextAccessor"/> to be registered in the service collection.
/// Ensure that <c>services.AddHttpContextAccessor()</c> has been called before using this method.
/// </para>
/// <para>
/// When <paramref name="options"/> is not supplied, the isolation key is derived from the
/// <see cref="ClaimTypes.NameIdentifier"/> claim, a stable unique subject identifier. For OpenID
/// Connect tokens (including Microsoft Entra ID), this is typically mapped from the <c>sub</c> claim
/// by the default JWT inbound claim mapping. Authentication schemes that do not project a unique
/// identifier onto <see cref="ClaimTypes.NameIdentifier"/> (or hosts that require a different claim
/// such as Entra's <c>oid</c>) should override
/// <see cref="ClaimsIdentitySessionIsolationKeyProviderOptions.ClaimType"/>; otherwise the key may be
/// absent, which causes strict-mode session stores to fail.
/// </para>
/// <para>
/// <strong>Security warning:</strong> If you override
/// <see cref="ClaimsIdentitySessionIsolationKeyProviderOptions.ClaimType"/>, the chosen claim must
/// uniquely identify the principal within the served population. Display names, usernames, email
/// aliases, and other mutable or non-unique claims are <strong>unsafe</strong> isolation keys unless
/// the host can prove their uniqueness across all callers, because distinct principals that share the
/// same claim value would receive the same isolation key and could access one another's sessions.
/// </para>
/// </remarks>
public static IServiceCollection UseClaimsBasedSessionIsolation(
this IServiceCollection services,
@@ -49,7 +49,7 @@ internal sealed class ForeachExecutor : DeclarativeActionExecutor<Foreach>
EvaluationResult<DataValue> expressionResult = this.Evaluator.GetValue(this.Model.Items);
if (expressionResult.Value is TableDataValue tableValue)
{
this._values = [.. tableValue.Values.Select(value => value.Properties.Values.First().ToFormula())];
this._values = [.. tableValue.Values.Select(ToLoopValue)];
}
else
{
@@ -99,6 +99,15 @@ internal sealed class ForeachExecutor : DeclarativeActionExecutor<Foreach>
}
}
// Power Fx wraps scalar array literals (`=[1, 2, 3]`) as `Table({Value: 1}, ...)`. Unwrap that single-column
// `Value`-record shape so `Local.LoopValue` is the scalar; multi-field and other shapes pass through unchanged.
private static FormulaValue ToLoopValue(DataValue value) =>
value is RecordDataValue record
&& record.Properties.Count == 1
&& record.Properties.TryGetValue("Value", out DataValue? singleColumn)
? singleColumn.ToFormula()
: value.ToFormula();
/// <inheritdoc/>
/// <remarks>
/// Persists the iteration cursor (<see cref="_index"/>), the materialized item snapshot
@@ -181,6 +181,36 @@ public sealed class ChatClientAgentOptions
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public bool EnableMessageInjection { get; set; }
/// <summary>
/// Gets or sets a value indicating whether to store automatically approved function calls in the session state
/// for tools that do not require approval when they are returned alongside tools that do.
/// </summary>
/// <remarks>
/// <para>
/// <see cref="FunctionInvokingChatClient"/> has an all-or-nothing behavior for approvals: when any tool
/// in a response is an <see cref="ApprovalRequiredAIFunction"/>, it converts all <see cref="FunctionCallContent"/>
/// items to <see cref="ToolApprovalRequestContent"/>, even for tools that do not require approval.
/// </para>
/// <para>
/// Setting this property to <see langword="true"/> injects an <see cref="NonApprovalRequiredFunctionBypassingChatClient"/>
/// decorator above <see cref="FunctionInvokingChatClient"/> in the pipeline. This decorator identifies approval
/// requests for non-approval-required tools, removes them from the response, and stores them in the session.
/// On the next request, the stored items are automatically re-injected as approved, so the caller only needs
/// to handle approval requests for tools that truly require human approval.
/// </para>
/// <para>
/// This option has no effect when <see cref="UseProvidedChatClientAsIs"/> is <see langword="true"/>.
/// When using a custom chat client stack, you can add an <see cref="NonApprovalRequiredFunctionBypassingChatClient"/>
/// manually via the <see cref="ChatClientBuilderExtensions.UseNonApprovalRequiredFunctionBypassing"/>
/// extension method.
/// </para>
/// </remarks>
/// <value>
/// Default is <see langword="false"/>.
/// </value>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public bool EnableNonApprovalRequiredFunctionBypassing { get; set; }
/// <summary>
/// Creates a new instance of <see cref="ChatClientAgentOptions"/> with the same values as this instance.
/// </summary>
@@ -199,5 +229,6 @@ public sealed class ChatClientAgentOptions
ThrowOnChatHistoryProviderConflict = this.ThrowOnChatHistoryProviderConflict,
RequirePerServiceCallChatHistoryPersistence = this.RequirePerServiceCallChatHistoryPersistence,
EnableMessageInjection = this.EnableMessageInjection,
EnableNonApprovalRequiredFunctionBypassing = this.EnableNonApprovalRequiredFunctionBypassing,
};
}
@@ -148,4 +148,35 @@ public static class ChatClientBuilderExtensions
{
return builder.Use(innerClient => new MessageInjectingChatClient(innerClient));
}
/// <summary>
/// Adds an <see cref="NonApprovalRequiredFunctionBypassingChatClient"/> to the chat client pipeline.
/// </summary>
/// <remarks>
/// <para>
/// This decorator should be positioned above the <see cref="FunctionInvokingChatClient"/> in the pipeline
/// so that it can intercept approval requests for tools that do not require approval. When
/// <see cref="FunctionInvokingChatClient"/> converts all function calls to approval requests (because at
/// least one tool requires approval), this decorator removes the requests for non-approval-required tools,
/// stores them in the session, and automatically re-injects them as approved on the next request.
/// </para>
/// <para>
/// This extension method is intended for use with custom chat client stacks when
/// <see cref="ChatClientAgentOptions.UseProvidedChatClientAsIs"/> is <see langword="true"/>.
/// When <see cref="ChatClientAgentOptions.UseProvidedChatClientAsIs"/> is <see langword="false"/> (the default),
/// the <see cref="ChatClientAgent"/> automatically injects this decorator when
/// <see cref="ChatClientAgentOptions.EnableNonApprovalRequiredFunctionBypassing"/> is <see langword="true"/>.
/// </para>
/// <para>
/// This decorator only works within the context of a running <see cref="ChatClientAgent"/> with
/// an active session, and will throw an exception if used in any other stack.
/// </para>
/// </remarks>
/// <param name="builder">The <see cref="ChatClientBuilder"/> to add the decorator to.</param>
/// <returns>The <paramref name="builder"/> for chaining.</returns>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public static ChatClientBuilder UseNonApprovalRequiredFunctionBypassing(this ChatClientBuilder builder)
{
return builder.Use(innerClient => new NonApprovalRequiredFunctionBypassingChatClient(innerClient));
}
}
@@ -53,6 +53,17 @@ public static class ChatClientExtensions
{
var chatBuilder = chatClient.AsBuilder();
// NonApprovalRequiredFunctionBypassingChatClient is registered before FunctionInvokingChatClient so that
// it sits above FICC in the pipeline. ChatClientBuilder.Build applies factories in reverse order,
// making the first Use() call outermost. By adding this decorator first, the resulting pipeline is:
// NonApprovalRequiredFunctionBypassingChatClient → FunctionInvokingChatClient → ChatHistoryPersistingChatClient → leaf IChatClient
// This allows the decorator to intercept FICC's responses and remove approval requests for tools
// that don't actually require approval, storing them for automatic re-injection on the next request.
if (options?.EnableNonApprovalRequiredFunctionBypassing is true)
{
chatBuilder.Use(innerClient => new NonApprovalRequiredFunctionBypassingChatClient(innerClient));
}
if (chatClient.GetService<FunctionInvokingChatClient>() is null)
{
chatBuilder.Use((innerClient, services) =>
@@ -0,0 +1,285 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
namespace Microsoft.Agents.AI;
/// <summary>
/// A delegating chat client that automatically removes <see cref="ToolApprovalRequestContent"/> for tools
/// that do not actually require approval, storing auto-approved results in the session for transparent
/// re-injection on the next request.
/// </summary>
/// <remarks>
/// <para>
/// <see cref="FunctionInvokingChatClient"/> has an all-or-nothing behavior for approvals: when any tool
/// in a response is an <see cref="ApprovalRequiredAIFunction"/>, it converts all <see cref="FunctionCallContent"/>
/// items to <see cref="ToolApprovalRequestContent"/> — even for tools that do not require approval. This
/// decorator sits above <see cref="FunctionInvokingChatClient"/> in the pipeline and transparently handles
/// the non-approval-required items so callers only see approval requests for tools that truly need them.
/// </para>
/// <para>
/// On outbound responses, the decorator identifies <see cref="ToolApprovalRequestContent"/> items for tools
/// that are not wrapped in <see cref="ApprovalRequiredAIFunction"/>, removes them from the response, and
/// stores them in the session's <see cref="AgentSessionStateBag"/>. On the next inbound request, the stored
/// items are re-injected as pre-approved <see cref="ToolApprovalResponseContent"/> so that
/// <see cref="FunctionInvokingChatClient"/> can process them alongside the caller's human-approved responses.
/// </para>
/// <para>
/// This decorator requires an active <see cref="AIAgent.CurrentRunContext"/> with a non-null
/// <see cref="AgentRunContext.Session"/>. An <see cref="InvalidOperationException"/> is thrown if no
/// run context or session is available.
/// </para>
/// </remarks>
internal sealed class NonApprovalRequiredFunctionBypassingChatClient : DelegatingChatClient
{
/// <summary>
/// The key used in <see cref="AgentSessionStateBag"/> to store pending auto-approved function calls
/// between agent runs.
/// </summary>
internal const string StateBagKey = "_autoApprovedFunctionCalls";
/// <summary>
/// Initializes a new instance of the <see cref="NonApprovalRequiredFunctionBypassingChatClient"/> class.
/// </summary>
/// <param name="innerClient">The underlying chat client (typically a <see cref="FunctionInvokingChatClient"/>).</param>
public NonApprovalRequiredFunctionBypassingChatClient(IChatClient innerClient)
: base(innerClient)
{
}
/// <inheritdoc/>
public override async Task<ChatResponse> GetResponseAsync(
IEnumerable<ChatMessage> messages,
ChatOptions? options = null,
CancellationToken cancellationToken = default)
{
var session = GetRequiredSession();
var autoApprovableNames = this.GetAutoApprovableToolNames(options);
messages = InjectPendingAutoApprovals(messages, session);
var response = await base.GetResponseAsync(messages, options, cancellationToken).ConfigureAwait(false);
RemoveAutoApprovedFromMessages(response.Messages, autoApprovableNames, session);
return response;
}
/// <inheritdoc/>
public override async IAsyncEnumerable<ChatResponseUpdate> GetStreamingResponseAsync(
IEnumerable<ChatMessage> messages,
ChatOptions? options = null,
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
var session = GetRequiredSession();
var autoApprovableNames = this.GetAutoApprovableToolNames(options);
messages = InjectPendingAutoApprovals(messages, session);
List<ToolApprovalRequestContent>? autoApproved = null;
try
{
await foreach (var update in base.GetStreamingResponseAsync(messages, options, cancellationToken).ConfigureAwait(false))
{
if (FilterUpdateContents(update, autoApprovableNames, ref autoApproved))
{
yield return update;
}
}
}
finally
{
if (autoApproved is { Count: > 0 })
{
session.StateBag.SetValue(StateBagKey, autoApproved, AgentJsonUtilities.DefaultOptions);
}
}
}
/// <summary>
/// Gets the current <see cref="AgentSession"/> from the ambient run context.
/// </summary>
/// <exception cref="InvalidOperationException">No run context or session is available.</exception>
private static AgentSession GetRequiredSession()
{
var runContext = AIAgent.CurrentRunContext
?? throw new InvalidOperationException(
$"{nameof(NonApprovalRequiredFunctionBypassingChatClient)} can only be used within the context of a running AIAgent. " +
"Ensure that the chat client is being invoked as part of an AIAgent.RunAsync or AIAgent.RunStreamingAsync call.");
return runContext.Session
?? throw new InvalidOperationException(
$"{nameof(NonApprovalRequiredFunctionBypassingChatClient)} requires a session. " +
"Ensure the agent has a resolved session before invoking the chat client.");
}
/// <summary>
/// Checks the session for stored auto-approvals from a previous turn and injects them as
/// a user message containing <see cref="ToolApprovalResponseContent"/> items appended to the input messages.
/// </summary>
/// <remarks>
/// All stored requests are unconditionally injected as approved responses regardless of whether the
/// tool set has changed, because the LLM requires a complete set of tool call responses for a prior turn.
/// </remarks>
private static IEnumerable<ChatMessage> InjectPendingAutoApprovals(
IEnumerable<ChatMessage> messages,
AgentSession session)
{
if (!session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
StateBagKey,
out var pendingRequests,
AgentJsonUtilities.DefaultOptions)
|| pendingRequests is not { Count: > 0 })
{
return messages;
}
session.StateBag.TryRemoveValue(StateBagKey);
List<AIContent> approvalResponses = [];
foreach (var request in pendingRequests)
{
approvalResponses.Add(request.CreateResponse(approved: true));
}
var userMessage = new ChatMessage(ChatRole.User, approvalResponses);
return messages.Concat([userMessage]);
}
/// <summary>
/// Builds a set of tool names that do not require approval and can be auto-approved,
/// by checking all available tools from <see cref="ChatOptions.Tools"/> and
/// <see cref="FunctionInvokingChatClient.AdditionalTools"/>.
/// </summary>
private HashSet<string> GetAutoApprovableToolNames(ChatOptions? options)
{
var ficc = this.GetService<FunctionInvokingChatClient>();
var allTools = (options?.Tools ?? Enumerable.Empty<AITool>())
.Concat(ficc?.AdditionalTools ?? Enumerable.Empty<AITool>());
return new HashSet<string>(
allTools
.OfType<AIFunction>()
.Where(static f => f.GetService<ApprovalRequiredAIFunction>() is null)
.Select(static f => f.Name),
StringComparer.Ordinal);
}
/// <summary>
/// Determines whether a <see cref="ToolApprovalRequestContent"/> can be auto-approved because
/// the underlying tool is not an <see cref="ApprovalRequiredAIFunction"/>.
/// </summary>
/// <returns>
/// <see langword="true"/> if the approval request is for a known tool that does not require approval
/// and can be auto-approved; <see langword="false"/> otherwise.
/// </returns>
private static bool IsAutoApprovable(ToolApprovalRequestContent approval, HashSet<string> autoApprovableNames)
{
if (approval.ToolCall is not FunctionCallContent fcc)
{
// Non-function tool calls cannot be auto-approved.
return false;
}
// Auto-approve only if the tool is known and explicitly does NOT require approval.
// Unknown tools are not in the set and are treated as approval-required (safe default).
return autoApprovableNames.Contains(fcc.Name);
}
/// <summary>
/// Scans response messages for auto-approvable <see cref="ToolApprovalRequestContent"/> items,
/// removes them from the messages, and stores them in the session for the next request.
/// </summary>
private static void RemoveAutoApprovedFromMessages(
IList<ChatMessage> messages,
HashSet<string> autoApprovableNames,
AgentSession session)
{
List<ToolApprovalRequestContent>? autoApproved = null;
foreach (var message in messages)
{
for (int i = message.Contents.Count - 1; i >= 0; i--)
{
if (message.Contents[i] is ToolApprovalRequestContent approval
&& IsAutoApprovable(approval, autoApprovableNames))
{
(autoApproved ??= []).Add(approval);
message.Contents.RemoveAt(i);
}
}
}
// Remove messages that are now empty after filtering.
for (int i = messages.Count - 1; i >= 0; i--)
{
if (messages[i].Contents.Count == 0)
{
messages.RemoveAt(i);
}
}
if (autoApproved is { Count: > 0 })
{
session.StateBag.SetValue(StateBagKey, autoApproved, AgentJsonUtilities.DefaultOptions);
}
}
/// <summary>
/// Filters auto-approvable <see cref="ToolApprovalRequestContent"/> items from a streaming update's
/// contents, collecting them for later storage.
/// </summary>
/// <returns>
/// <see langword="true"/> if the update should be yielded (has remaining content or had no
/// approval content to begin with); <see langword="false"/> if the update is now empty and
/// should be skipped.
/// </returns>
private static bool FilterUpdateContents(
ChatResponseUpdate update,
HashSet<string> autoApprovableNames,
ref List<ToolApprovalRequestContent>? autoApproved)
{
bool hasApprovalContent = false;
List<AIContent> filteredContents = [];
bool removedAny = false;
for (int i = 0; i < update.Contents.Count; i++)
{
var content = update.Contents[i];
if (content is ToolApprovalRequestContent approval)
{
hasApprovalContent = true;
if (IsAutoApprovable(approval, autoApprovableNames))
{
(autoApproved ??= []).Add(approval);
removedAny = true;
}
else
{
filteredContents.Add(content);
}
}
else
{
filteredContents.Add(content);
}
}
if (removedAny)
{
update.Contents = filteredContents;
}
// Yield the update unless it was purely auto-approvable approval content (now empty).
return update.Contents.Count > 0 || !hasApprovalContent;
}
}
@@ -51,20 +51,22 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
{
private readonly ProviderSessionState<ToolApprovalState> _sessionState;
private readonly JsonSerializerOptions _jsonSerializerOptions;
private readonly Func<FunctionCallContent, ValueTask<bool>>[]? _autoApprovalRules;
/// <summary>
/// Initializes a new instance of the <see cref="ToolApprovalAgent"/> class.
/// </summary>
/// <param name="innerAgent">The underlying agent to delegate to.</param>
/// <param name="jsonSerializerOptions">
/// Optional <see cref="JsonSerializerOptions"/> used for serializing argument values when storing rules
/// and for persisting state. When <see langword="null"/>, <see cref="AgentJsonUtilities.DefaultOptions"/> is used.
/// <param name="options">
/// Optional <see cref="ToolApprovalAgentOptions"/> for configuring serialization and auto-approval rules.
/// When <see langword="null"/>, default settings are used.
/// </param>
/// <exception cref="ArgumentNullException"><paramref name="innerAgent"/> is <see langword="null"/>.</exception>
public ToolApprovalAgent(AIAgent innerAgent, JsonSerializerOptions? jsonSerializerOptions = null)
public ToolApprovalAgent(AIAgent innerAgent, ToolApprovalAgentOptions? options = null)
: base(innerAgent)
{
this._jsonSerializerOptions = jsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
this._jsonSerializerOptions = options?.JsonSerializerOptions ?? AgentJsonUtilities.DefaultOptions;
this._autoApprovalRules = options?.AutoApprovalRules?.ToArray();
this._sessionState = new ProviderSessionState<ToolApprovalState>(
_ => new ToolApprovalState(),
"toolApprovalState",
@@ -79,7 +81,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
CancellationToken cancellationToken = default)
{
// Steps 12: Unwrap AlwaysApprove wrappers, process any queued approval requests.
var (state, callerMessages, nextQueuedItem) = this.PrepareInboundMessages(messages, session);
var (state, callerMessages, nextQueuedItem) = await this.PrepareInboundMessagesAsync(messages, session).ConfigureAwait(false);
if (nextQueuedItem is not null)
{
@@ -98,7 +100,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
var response = await this.InnerAgent.RunAsync(processedMessages, session, options, cancellationToken).ConfigureAwait(false);
// Classify approval requests: auto-approve matching, queue excess, keep first unapproved.
bool allAutoApproved = this.ProcessAndQueueOutboundApprovalRequests(response.Messages, state, session);
bool allAutoApproved = await this.ProcessAndQueueOutboundApprovalRequestsAsync(response.Messages, state, session).ConfigureAwait(false);
if (!allAutoApproved)
{
@@ -119,7 +121,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
[EnumeratorCancellation] CancellationToken cancellationToken = default)
{
// Steps 12: Unwrap AlwaysApprove wrappers, process any queued approval requests.
var (state, callerMessages, nextQueuedItem) = this.PrepareInboundMessages(messages, session);
var (state, callerMessages, nextQueuedItem) = await this.PrepareInboundMessagesAsync(messages, session).ConfigureAwait(false);
if (nextQueuedItem is not null)
{
@@ -197,7 +199,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
yield break;
}
// 4. Classify the collected approval requests against standing rules.
// 4. Classify the collected approval requests against standing rules and auto-approval rules.
List<ToolApprovalRequestContent> unapproved = [];
foreach (var tarc in streamedApprovalRequests)
{
@@ -206,6 +208,11 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
}
else if (await this.MatchesAutoApprovalRuleAsync(tarc).ConfigureAwait(false))
{
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by auto-approval rule"));
}
else
{
unapproved.Add(tarc);
@@ -291,9 +298,9 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
}
/// <summary>
/// Re-evaluates queued approval requests against current rules and auto-approves any that now match.
/// Re-evaluates queued approval requests against current rules and auto-approval rules, and auto-approves any that now match.
/// </summary>
private void DrainAutoApprovableFromQueue(ToolApprovalState state)
private async ValueTask DrainAutoApprovableFromQueueAsync(ToolApprovalState state)
{
for (int i = state.QueuedApprovalRequests.Count - 1; i >= 0; i--)
{
@@ -303,6 +310,12 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
state.QueuedApprovalRequests[i].CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
state.QueuedApprovalRequests.RemoveAt(i);
}
else if (await this.MatchesAutoApprovalRuleAsync(state.QueuedApprovalRequests[i]).ConfigureAwait(false))
{
state.CollectedApprovalResponses.Add(
state.QueuedApprovalRequests[i].CreateResponse(approved: true, reason: "Auto-approved by auto-approval rule"));
state.QueuedApprovalRequests.RemoveAt(i);
}
}
}
@@ -318,8 +331,8 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
/// A tuple of (state, processed caller messages, next queued item or <see langword="null"/> if the queue is resolved).
/// When the returned item is non-null, the caller should return/yield it without calling the inner agent.
/// </returns>
private (ToolApprovalState State, List<ChatMessage> CallerMessages, ToolApprovalRequestContent? NextQueuedItem)
PrepareInboundMessages(IEnumerable<ChatMessage> messages, AgentSession? session)
private async ValueTask<(ToolApprovalState State, List<ChatMessage> CallerMessages, ToolApprovalRequestContent? NextQueuedItem)>
PrepareInboundMessagesAsync(IEnumerable<ChatMessage> messages, AgentSession? session)
{
var state = this._sessionState.GetOrInitializeState(session);
@@ -337,7 +350,7 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
// Re-evaluate remaining queued items — the caller may have added new rules
// (e.g., "always approve this tool") that resolve additional items.
this.DrainAutoApprovableFromQueue(state);
await this.DrainAutoApprovableFromQueueAsync(state).ConfigureAwait(false);
if (state.QueuedApprovalRequests.Count > 0)
{
@@ -386,15 +399,18 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
/// <see langword="true"/> if all TARc items were auto-approved (caller should re-invoke the inner agent);
/// <see langword="false"/> otherwise.
/// </returns>
private bool ProcessAndQueueOutboundApprovalRequests(
private async ValueTask<bool> ProcessAndQueueOutboundApprovalRequestsAsync(
IList<ChatMessage> responseMessages,
ToolApprovalState state,
AgentSession? session)
{
// Pass 1: Scan all response messages and classify each approval request as
// auto-approved (matches a standing rule) or unapproved (needs caller decision).
var autoApproved = new List<ToolApprovalRequestContent>();
// Pass 1: Scan all response messages and classify each approval request.
// Auto-approved requests (matching a standing rule or auto-approval rule) have their
// responses collected immediately, preserving the original request order, and are
// marked for removal. Unapproved requests are collected for the caller to decide.
var toRemove = new HashSet<ToolApprovalRequestContent>();
var unapproved = new List<ToolApprovalRequestContent>();
int autoApprovedCount = 0;
foreach (var message in responseMessages)
{
@@ -404,7 +420,17 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
{
if (MatchesRule(tarc, state.Rules, this._jsonSerializerOptions))
{
autoApproved.Add(tarc);
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
toRemove.Add(tarc);
autoApprovedCount++;
}
else if (await this.MatchesAutoApprovalRuleAsync(tarc).ConfigureAwait(false))
{
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by auto-approval rule"));
toRemove.Add(tarc);
autoApprovedCount++;
}
else
{
@@ -415,18 +441,12 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
}
// Nothing to process: no auto-approved items and at most one unapproved (no queueing needed).
if (autoApproved.Count == 0 && unapproved.Count <= 1)
// No responses were collected above in this case, so state is unmodified and safe to leave.
if (autoApprovedCount == 0 && unapproved.Count <= 1)
{
return false;
}
// Store auto-approved responses for later injection into the inner agent.
foreach (var tarc in autoApproved)
{
state.CollectedApprovalResponses.Add(
tarc.CreateResponse(approved: true, reason: "Auto-approved by standing rule"));
}
// If every approval request was auto-approved, strip them all and signal the caller
// to re-invoke the inner agent immediately with the collected responses.
if (unapproved.Count == 0)
@@ -439,14 +459,10 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
// Pass 2: Keep only the first unapproved request in the response (for the caller to decide).
// Queue the remaining unapproved requests for subsequent one-at-a-time delivery.
// Remove all auto-approved and queued items from the response messages.
var toRemove = new HashSet<ToolApprovalRequestContent>(autoApproved);
if (unapproved.Count > 1)
for (int i = 1; i < unapproved.Count; i++)
{
for (int i = 1; i < unapproved.Count; i++)
{
toRemove.Add(unapproved[i]);
state.QueuedApprovalRequests.Add(unapproved[i]);
}
toRemove.Add(unapproved[i]);
state.QueuedApprovalRequests.Add(unapproved[i]);
}
// Walk messages in reverse and strip marked items.
@@ -663,8 +679,36 @@ public sealed class ToolApprovalAgent : DelegatingAIAgent
}
/// <summary>
/// Compares stored rule arguments against actual function call arguments for an exact match.
/// Checks whether a <see cref="ToolApprovalRequestContent"/> is approved by any of the configured
/// auto-approval rules (heuristic functions).
/// </summary>
/// <returns>
/// <see langword="true"/> if any auto-approval rule returns <see langword="true"/> for the function call;
/// <see langword="false"/> if no rules are configured, the request is not a function call, or no rule approves it.
/// </returns>
private async ValueTask<bool> MatchesAutoApprovalRuleAsync(ToolApprovalRequestContent request)
{
if (this._autoApprovalRules is not { Length: > 0 })
{
return false;
}
if (request.ToolCall is not FunctionCallContent functionCall)
{
return false;
}
foreach (var rule in this._autoApprovalRules)
{
if (await rule(functionCall).ConfigureAwait(false))
{
return true;
}
}
return false;
}
private static bool ArgumentsMatch(IDictionary<string, string> ruleArguments, IDictionary<string, object?>? callArguments, JsonSerializerOptions jsonSerializerOptions)
{
if (callArguments is null)
@@ -1,7 +1,6 @@
// Copyright (c) Microsoft. All rights reserved.
using System.Diagnostics.CodeAnalysis;
using System.Text.Json;
using Microsoft.Shared.DiagnosticIds;
using Microsoft.Shared.Diagnostics;
@@ -17,9 +16,9 @@ public static class ToolApprovalAgentBuilderExtensions
/// Adds tool approval middleware to the agent pipeline, enabling "don't ask again" approval behavior.
/// </summary>
/// <param name="builder">The <see cref="AIAgentBuilder"/> to which tool approval support will be added.</param>
/// <param name="jsonSerializerOptions">
/// Optional <see cref="JsonSerializerOptions"/> used for serializing argument values when storing rules
/// and for persisting state. When <see langword="null"/>, <see cref="AgentJsonUtilities.DefaultOptions"/> is used.
/// <param name="options">
/// Optional <see cref="ToolApprovalAgentOptions"/> for configuring serialization and auto-approval rules.
/// When <see langword="null"/>, default settings are used.
/// </param>
/// <returns>The <see cref="AIAgentBuilder"/> with tool approval middleware added, enabling method chaining.</returns>
/// <exception cref="System.ArgumentNullException"><paramref name="builder"/> is <see langword="null"/>.</exception>
@@ -32,6 +31,6 @@ public static class ToolApprovalAgentBuilderExtensions
/// </remarks>
public static AIAgentBuilder UseToolApproval(
this AIAgentBuilder builder,
JsonSerializerOptions? jsonSerializerOptions = null)
=> Throw.IfNull(builder).Use(innerAgent => new ToolApprovalAgent(innerAgent, jsonSerializerOptions));
ToolApprovalAgentOptions? options = null)
=> Throw.IfNull(builder).Use(innerAgent => new ToolApprovalAgent(innerAgent, options));
}
@@ -0,0 +1,45 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Diagnostics.CodeAnalysis;
using System.Text.Json;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Microsoft.Shared.DiagnosticIds;
namespace Microsoft.Agents.AI;
/// <summary>
/// Options for configuring the <see cref="ToolApprovalAgent"/> middleware.
/// </summary>
[Experimental(DiagnosticIds.Experiments.AgentsAIExperiments)]
public class ToolApprovalAgentOptions
{
/// <summary>
/// Gets or sets the <see cref="System.Text.Json.JsonSerializerOptions"/> used for serializing argument values
/// when storing rules and for persisting state.
/// </summary>
/// <remarks>
/// When <see langword="null"/>, <see cref="AgentJsonUtilities.DefaultOptions"/> is used.
/// </remarks>
public JsonSerializerOptions? JsonSerializerOptions { get; set; }
/// <summary>
/// Gets or sets a collection of heuristic functions that can automatically approve function calls
/// that would otherwise require user approval.
/// </summary>
/// <remarks>
/// <para>
/// Each function receives a <see cref="FunctionCallContent"/> representing the tool call that requires approval
/// and returns a <see cref="ValueTask{Boolean}"/> that resolves to <see langword="true"/> to auto-approve
/// the call, or <see langword="false"/> to continue evaluating the next rule.
/// </para>
/// <para>
/// Auto-approval rules are evaluated after standing rules (derived from prior user approvals) but before
/// prompting the user. Rules are evaluated in order; the first rule returning <see langword="true"/>
/// causes the function call to be auto-approved.
/// </para>
/// </remarks>
public IEnumerable<Func<FunctionCallContent, ValueTask<bool>>>? AutoApprovalRules { get; set; }
}
@@ -182,7 +182,7 @@ public sealed class WorkflowConsoleAppSamplesValidation(ITestOutputHelper output
}
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "KeyNotFoundException in workflow execution. See https://github.com/microsoft/agent-framework/issues/6404")]
public async Task WorkflowEventsSampleValidationAsync()
{
using CancellationTokenSource testTimeoutCts = this.CreateTestTimeoutCts(s_testTimeout);
@@ -278,7 +278,7 @@ public sealed class WorkflowConsoleAppSamplesValidation(ITestOutputHelper output
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "KeyNotFoundException in workflow execution. See https://github.com/microsoft/agent-framework/issues/6404")]
public async Task WorkflowSharedStateSampleValidationAsync()
{
using CancellationTokenSource testTimeoutCts = this.CreateTestTimeoutCts(s_testTimeout);
@@ -376,7 +376,7 @@ public sealed class WorkflowConsoleAppSamplesValidation(ITestOutputHelper output
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "KeyNotFoundException in workflow execution. See https://github.com/microsoft/agent-framework/issues/6404")]
public async Task SubWorkflowsSampleValidationAsync()
{
using CancellationTokenSource testTimeoutCts = this.CreateTestTimeoutCts(s_testTimeout);
@@ -452,7 +452,7 @@ public sealed class WorkflowConsoleAppSamplesValidation(ITestOutputHelper output
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "KeyNotFoundException in workflow execution. See https://github.com/microsoft/agent-framework/issues/6404")]
public async Task WorkflowHITLSampleValidationAsync()
{
using CancellationTokenSource testTimeoutCts = this.CreateTestTimeoutCts(s_testTimeout);
@@ -43,7 +43,7 @@ public class HostedFoundryMemoryProviderScopesTests
}
[Fact]
public void PerUserAndChat_ComposesUserAndChatWithColon()
public void PerUserAndChat_ComposesUserAndChatWithEscapedSeparator()
{
// Arrange
var session = CreateTaggedSession(TestUserId, TestChatId);
@@ -54,7 +54,51 @@ public class HostedFoundryMemoryProviderScopesTests
// Assert
Assert.NotNull(state);
Assert.Equal($"{TestUserId}:{TestChatId}", state.Scope.Scope);
Assert.Equal($"{TestUserId}::{TestChatId}", state.Scope.Scope);
}
[Fact]
public void PerUserAndChat_EscapesColonsInUserAndChat()
{
// Arrange
var session = CreateTaggedSession("alice:finance", "q2:final");
var initializer = HostedFoundryMemoryProviderScopes.PerUserAndChat();
// Act
var state = initializer(session);
// Assert - colons inside each part are escaped as \: , parts joined with ::
Assert.Equal(@"alice\:finance::q2\:final", state.Scope.Scope);
}
[Fact]
public void PerUserAndChat_EscapesBackslashesInUserAndChat()
{
// Arrange
var session = CreateTaggedSession(@"alice\corp", @"chat\1");
var initializer = HostedFoundryMemoryProviderScopes.PerUserAndChat();
// Act
var state = initializer(session);
// Assert - backslashes escaped first as \\ , parts joined with ::
Assert.Equal(@"alice\\corp::chat\\1", state.Scope.Scope);
}
[Fact]
public void PerUserAndChat_DistinctContextsDoNotCollide()
{
// Arrange - two distinct (UserId, ChatId) pairs that collide under raw-colon composition.
var sessionA = CreateTaggedSession("alice:finance", "q2");
var sessionB = CreateTaggedSession("alice", "finance:q2");
var initializer = HostedFoundryMemoryProviderScopes.PerUserAndChat();
// Act
var scopeA = initializer(sessionA).Scope.Scope;
var scopeB = initializer(sessionB).Scope.Scope;
// Assert
Assert.NotEqual(scopeA, scopeB);
}
[Fact]
@@ -4,7 +4,8 @@ using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading.Tasks;
using GitHub.Copilot.SDK;
using GitHub.Copilot;
using GitHub.Copilot.Rpc;
using Microsoft.Extensions.AI;
namespace Microsoft.Agents.AI.GitHub.Copilot.IntegrationTests;
@@ -13,8 +14,8 @@ public class GitHubCopilotAgentTests
{
private const string SkipReason = "Integration tests require GitHub Copilot CLI installed. For local execution only.";
private static Task<PermissionRequestResult> OnPermissionRequestAsync(PermissionRequest request, PermissionInvocation invocation)
=> Task.FromResult(new PermissionRequestResult { Kind = PermissionRequestResultKind.Approved });
private static Task<PermissionDecision> OnPermissionRequestAsync(PermissionRequest request, PermissionInvocation invocation)
=> Task.FromResult(PermissionDecision.ApproveOnce());
[Fact(Skip = SkipReason)]
public async Task RunAsync_WithSimplePrompt_ReturnsResponseAsync()
@@ -3,6 +3,7 @@
<PropertyGroup>
<!-- GitHub.Copilot.SDK only supports .NET 8.0+ -->
<TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
<NoWarn>$(NoWarn);GHCP001</NoWarn>
</PropertyGroup>
<ItemGroup>
@@ -2,7 +2,7 @@
using System;
using System.Collections.Generic;
using GitHub.Copilot.SDK;
using GitHub.Copilot;
using Microsoft.Extensions.AI;
namespace Microsoft.Agents.AI.GitHub.Copilot.UnitTests;
@@ -16,7 +16,7 @@ public sealed class CopilotClientExtensionsTests
public void AsAIAgent_WithAllParameters_ReturnsGitHubCopilotAgentWithSpecifiedProperties()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
const string TestId = "test-agent-id";
const string TestName = "Test Agent";
@@ -37,7 +37,7 @@ public sealed class CopilotClientExtensionsTests
public void AsAIAgent_WithMinimalParameters_ReturnsGitHubCopilotAgent()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
// Act
var agent = copilotClient.AsAIAgent(ownsClient: false, tools: null);
@@ -61,7 +61,7 @@ public sealed class CopilotClientExtensionsTests
public void AsAIAgent_WithOwnsClient_ReturnsAgentThatOwnsClient()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
// Act
var agent = copilotClient.AsAIAgent(ownsClient: true, tools: null);
@@ -75,7 +75,7 @@ public sealed class CopilotClientExtensionsTests
public void AsAIAgent_WithTools_ReturnsAgentWithTools()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
List<AITool> tools = [AIFunctionFactory.Create(() => "test", "TestFunc", "Test function")];
// Act
@@ -3,7 +3,8 @@
using System;
using System.Collections.Generic;
using System.Threading.Tasks;
using GitHub.Copilot.SDK;
using GitHub.Copilot;
using GitHub.Copilot.Rpc;
using Microsoft.Extensions.AI;
namespace Microsoft.Agents.AI.GitHub.Copilot.UnitTests;
@@ -17,7 +18,7 @@ public sealed class GitHubCopilotAgentTests
public void Constructor_WithCopilotClient_InitializesPropertiesCorrectly()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
const string TestId = "test-id";
const string TestName = "test-name";
const string TestDescription = "test-description";
@@ -42,7 +43,7 @@ public sealed class GitHubCopilotAgentTests
public void Constructor_WithDefaultParameters_UsesBaseProperties()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
// Act
var agent = new GitHubCopilotAgent(copilotClient, ownsClient: false, tools: null);
@@ -58,7 +59,7 @@ public sealed class GitHubCopilotAgentTests
public async Task CreateSessionAsync_ReturnsGitHubCopilotAgentSessionAsync()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
var agent = new GitHubCopilotAgent(copilotClient, ownsClient: false, tools: null);
// Act
@@ -73,7 +74,7 @@ public sealed class GitHubCopilotAgentTests
public async Task CreateSessionAsync_WithSessionId_ReturnsSessionWithSessionIdAsync()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
var agent = new GitHubCopilotAgent(copilotClient, ownsClient: false, tools: null);
const string TestSessionId = "test-session-id";
@@ -90,7 +91,7 @@ public sealed class GitHubCopilotAgentTests
public void Constructor_WithTools_InitializesCorrectly()
{
// Arrange
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
List<AITool> tools = [AIFunctionFactory.Create(() => "test", "TestFunc", "Test function")];
// Act
@@ -105,12 +106,12 @@ public sealed class GitHubCopilotAgentTests
public void CopySessionConfig_CopiesAllProperties()
{
// Arrange
List<AIFunction> tools = [AIFunctionFactory.Create(() => "test", "TestFunc", "Test function")];
List<AIFunctionDeclaration> tools = [AIFunctionFactory.Create(() => "test", "TestFunc", "Test function")];
var hooks = new SessionHooks();
var infiniteSessions = new InfiniteSessionConfig();
var systemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = "Be helpful" };
PermissionRequestHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
UserInputHandler userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
Func<PermissionRequest, PermissionInvocation, Task<PermissionDecision>> permissionHandler = (_, _) => Task.FromResult(PermissionDecision.ApproveOnce());
Func<UserInputRequest, UserInputInvocation, Task<UserInputResponse>> userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
var mcpServers = new Dictionary<string, McpServerConfig> { ["server1"] = new McpStdioServerConfig() };
var source = new SessionConfig
@@ -122,7 +123,7 @@ public sealed class GitHubCopilotAgentTests
AvailableTools = ["tool1", "tool2"],
ExcludedTools = ["tool3"],
WorkingDirectory = "/workspace",
ConfigDir = "/config",
ConfigDirectory = "/config",
Hooks = hooks,
InfiniteSessions = infiniteSessions,
OnPermissionRequest = permissionHandler,
@@ -137,17 +138,15 @@ public sealed class GitHubCopilotAgentTests
// Assert
Assert.Equal("gpt-4o", result.Model);
Assert.Equal("high", result.ReasoningEffort);
Assert.Same(tools, result.Tools);
Assert.Same(systemMessage, result.SystemMessage);
Assert.Equal(systemMessage, result.SystemMessage);
Assert.Equal(new List<string> { "tool1", "tool2" }, result.AvailableTools);
Assert.Equal(new List<string> { "tool3" }, result.ExcludedTools);
Assert.Equal("/workspace", result.WorkingDirectory);
Assert.Equal("/config", result.ConfigDir);
Assert.Equal("/config", result.ConfigDirectory);
Assert.Same(hooks, result.Hooks);
Assert.Same(infiniteSessions, result.InfiniteSessions);
Assert.Same(permissionHandler, result.OnPermissionRequest);
Assert.Same(userInputHandler, result.OnUserInputRequest);
Assert.Same(mcpServers, result.McpServers);
Assert.Equal(new List<string> { "skill1" }, result.DisabledSkills);
Assert.True(result.Streaming);
}
@@ -156,12 +155,12 @@ public sealed class GitHubCopilotAgentTests
public void CopyResumeSessionConfig_CopiesAllProperties()
{
// Arrange
List<AIFunction> tools = [AIFunctionFactory.Create(() => "test", "TestFunc", "Test function")];
List<AIFunctionDeclaration> tools = [AIFunctionFactory.Create(() => "test", "TestFunc", "Test function")];
var hooks = new SessionHooks();
var infiniteSessions = new InfiniteSessionConfig();
var systemMessage = new SystemMessageConfig { Mode = SystemMessageMode.Append, Content = "Be helpful" };
PermissionRequestHandler permissionHandler = (_, _) => Task.FromResult(new PermissionRequestResult());
UserInputHandler userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
Func<PermissionRequest, PermissionInvocation, Task<PermissionDecision>> permissionHandler = (_, _) => Task.FromResult(PermissionDecision.ApproveOnce());
Func<UserInputRequest, UserInputInvocation, Task<UserInputResponse>> userInputHandler = (_, _) => Task.FromResult(new UserInputResponse { Answer = "input" });
var mcpServers = new Dictionary<string, McpServerConfig> { ["server1"] = new McpStdioServerConfig() };
var source = new SessionConfig
@@ -173,7 +172,7 @@ public sealed class GitHubCopilotAgentTests
AvailableTools = ["tool1", "tool2"],
ExcludedTools = ["tool3"],
WorkingDirectory = "/workspace",
ConfigDir = "/config",
ConfigDirectory = "/config",
Hooks = hooks,
InfiniteSessions = infiniteSessions,
OnPermissionRequest = permissionHandler,
@@ -193,7 +192,7 @@ public sealed class GitHubCopilotAgentTests
Assert.Equal(new List<string> { "tool1", "tool2" }, result.AvailableTools);
Assert.Equal(new List<string> { "tool3" }, result.ExcludedTools);
Assert.Equal("/workspace", result.WorkingDirectory);
Assert.Equal("/config", result.ConfigDir);
Assert.Equal("/config", result.ConfigDirectory);
Assert.Same(hooks, result.Hooks);
Assert.Same(infiniteSessions, result.InfiniteSessions);
Assert.Same(permissionHandler, result.OnPermissionRequest);
@@ -218,7 +217,7 @@ public sealed class GitHubCopilotAgentTests
Assert.Null(result.OnUserInputRequest);
Assert.Null(result.Hooks);
Assert.Null(result.WorkingDirectory);
Assert.Null(result.ConfigDir);
Assert.Null(result.ConfigDirectory);
Assert.True(result.Streaming);
}
@@ -233,7 +232,7 @@ public sealed class GitHubCopilotAgentTests
Content = "Some streamed content that was already delivered via delta events"
}
};
CopilotClient copilotClient = new(new CopilotClientOptions { AutoStart = false });
CopilotClient copilotClient = new(new CopilotClientOptions());
const string TestId = "agent-id";
var agent = new GitHubCopilotAgent(copilotClient, ownsClient: false, id: TestId, tools: null);
AgentResponseUpdate result = agent.ConvertToAgentResponseUpdate(assistantMessage);
@@ -3,6 +3,7 @@
<PropertyGroup>
<!-- GitHub.Copilot.SDK only supports .NET 8.0+ -->
<TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
<NoWarn>$(NoWarn);GHCP001</NoWarn>
</PropertyGroup>
<ItemGroup>
@@ -27,6 +27,7 @@ public class HarnessAgentOptionsTests
Assert.Null(options.ChatHistoryProvider);
Assert.Null(options.AIContextProviders);
Assert.False(options.DisableToolApproval);
Assert.False(options.DisableNonApprovalRequiredFunctionBypassing);
Assert.False(options.DisableFileMemory);
Assert.False(options.DisableFileAccess);
Assert.False(options.DisableWebSearch);
@@ -80,6 +81,7 @@ public class HarnessAgentOptionsTests
AIContextProviders = contextProviders,
MaximumIterationsPerRequest = 42,
DisableToolApproval = true,
DisableNonApprovalRequiredFunctionBypassing = true,
DisableFileMemory = true,
FileMemoryStore = fileMemoryStore,
DisableFileAccess = true,
@@ -112,6 +114,7 @@ public class HarnessAgentOptionsTests
Assert.Same(contextProviders, options.AIContextProviders);
Assert.Equal(42, options.MaximumIterationsPerRequest);
Assert.True(options.DisableToolApproval);
Assert.True(options.DisableNonApprovalRequiredFunctionBypassing);
Assert.True(options.DisableFileMemory);
Assert.Same(fileMemoryStore, options.FileMemoryStore);
Assert.True(options.DisableFileAccess);
@@ -644,6 +644,142 @@ public class HarnessAgentTests
Assert.Null(agent.GetService<ToolApprovalAgent>());
}
/// <summary>
/// Verify that ToolApprovalAgentOptions auto-approval rules are passed through and actually used.
/// </summary>
[Fact]
public async Task ToolApproval_AutoApprovalRulesAreAppliedAsync()
{
// Arrange — inner client returns an approval request on first call, then final response on second.
var callCount = 0;
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadTool"));
var mockClient = new Mock<IChatClient>();
mockClient
.Setup(c => c.GetResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(() =>
{
callCount++;
if (callCount == 1)
{
return new ChatResponse(new ChatMessage(ChatRole.Assistant, [approvalRequest]));
}
return new ChatResponse(new ChatMessage(ChatRole.Assistant, "Done"));
});
var options = CreateAllDisabledOptions();
options.DisableToolApproval = false;
options.ToolApprovalAgentOptions = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(fcc.Name == "ReadTool")]
};
var agent = new HarnessAgent(mockClient.Object, TestMaxContextWindowTokens, TestMaxOutputTokens, options);
var session = await agent.CreateSessionAsync();
// Act
var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session);
// Assert — the auto-approval rule approved the request, so we get "Done" (not an approval request)
Assert.Equal(2, callCount);
Assert.Equal("Done", response.Text);
}
#endregion
#region Feature: NonApprovalRequiredFunctionBypassing
/// <summary>
/// Verify that by default, when a response contains a mix of tools that require approval and tools that do not,
/// only the approval-required tool is surfaced to the caller. The non-approval-required tool is bypassed
/// (stored as auto-approved) by the <c>NonApprovalRequiredFunctionBypassingChatClient</c> decorator.
/// </summary>
[Fact]
public async Task NonApprovalRequiredFunctionBypassing_BypassesNonApprovalToolsByDefaultAsync()
{
// Arrange — the model requests both a normal tool and an approval-required tool in the same turn.
var normalTool = AIFunctionFactory.Create(() => "result", "NormalTool");
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "ApprovalTool"));
var mockClient = new Mock<IChatClient>();
mockClient
.Setup(c => c.GetResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(() => new ChatResponse(new ChatMessage(ChatRole.Assistant,
[
new FunctionCallContent("call1", "NormalTool"),
new FunctionCallContent("call2", "ApprovalTool"),
])));
// Disable ToolApproval so the approval requests surface in the response instead of being handled.
var options = CreateAllDisabledOptions();
options.ChatOptions = new ChatOptions { Tools = [normalTool, approvalTool] };
var agent = new HarnessAgent(mockClient.Object, TestMaxContextWindowTokens, TestMaxOutputTokens, options);
var session = await agent.CreateSessionAsync();
// Act
var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session);
// Assert — only the approval-required tool surfaces as an approval request; the normal tool is bypassed.
var approvalRequests = response.Messages
.SelectMany(m => m.Contents)
.OfType<ToolApprovalRequestContent>()
.ToList();
var approvalRequest = Assert.Single(approvalRequests);
Assert.Equal("ApprovalTool", Assert.IsType<FunctionCallContent>(approvalRequest.ToolCall).Name);
}
/// <summary>
/// Verify that when bypassing is disabled, all tools (including those that do not require approval) are surfaced
/// as approval requests, reflecting the all-or-nothing behavior of <see cref="FunctionInvokingChatClient"/>.
/// </summary>
[Fact]
public async Task NonApprovalRequiredFunctionBypassing_SurfacesAllApprovalsWhenDisabledAsync()
{
// Arrange — the model requests both a normal tool and an approval-required tool in the same turn.
var normalTool = AIFunctionFactory.Create(() => "result", "NormalTool");
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "ApprovalTool"));
var mockClient = new Mock<IChatClient>();
mockClient
.Setup(c => c.GetResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions>(),
It.IsAny<CancellationToken>()))
.ReturnsAsync(() => new ChatResponse(new ChatMessage(ChatRole.Assistant,
[
new FunctionCallContent("call1", "NormalTool"),
new FunctionCallContent("call2", "ApprovalTool"),
])));
var options = CreateAllDisabledOptions();
options.DisableNonApprovalRequiredFunctionBypassing = true;
options.ChatOptions = new ChatOptions { Tools = [normalTool, approvalTool] };
var agent = new HarnessAgent(mockClient.Object, TestMaxContextWindowTokens, TestMaxOutputTokens, options);
var session = await agent.CreateSessionAsync();
// Act
var response = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session);
// Assert — both tools surface as approval requests because bypassing is disabled.
var approvalRequests = response.Messages
.SelectMany(m => m.Contents)
.OfType<ToolApprovalRequestContent>()
.Select(r => ((FunctionCallContent)r.ToolCall).Name)
.ToList();
Assert.Equal(2, approvalRequests.Count);
Assert.Contains("NormalTool", approvalRequests);
Assert.Contains("ApprovalTool", approvalRequests);
}
#endregion
#region Feature: OpenTelemetry
@@ -60,7 +60,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
await Task.CompletedTask;
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task SingleAgentSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "01_SingleAgent");
@@ -148,7 +148,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task MultiAgentOrchestrationConcurrentSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "03_AgentOrchestration_Concurrency");
@@ -198,7 +198,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task MultiAgentOrchestrationConditionalsSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "04_AgentOrchestration_Conditionals");
@@ -216,7 +216,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task SingleAgentOrchestrationHITLSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "05_AgentOrchestration_HITL");
@@ -272,7 +272,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task LongRunningToolsSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "06_LongRunningTools");
@@ -362,7 +362,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task AgentAsMcpToolAsync()
{
string samplePath = Path.Combine(s_samplesPath, "07_AgentAsMcpTool");
@@ -402,7 +402,7 @@ public sealed class SamplesValidation(ITestOutputHelper outputHelper) : IAsyncLi
});
}
[RetryFact(2, 5000)]
[RetryFact(2, 5000, Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task ReliableStreamingSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "08_ReliableStreaming");
@@ -62,7 +62,7 @@ public sealed class WorkflowSamplesValidation(ITestOutputHelper outputHelper) :
return default;
}
[Fact]
[Fact(Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task SequentialWorkflowSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "01_SequentialWorkflow");
@@ -168,7 +168,7 @@ public sealed class WorkflowSamplesValidation(ITestOutputHelper outputHelper) :
});
}
[Fact]
[Fact(Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task HITLWorkflowSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "03_WorkflowHITL");
@@ -277,7 +277,7 @@ public sealed class WorkflowSamplesValidation(ITestOutputHelper outputHelper) :
});
}
[Fact]
[Fact(Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task WorkflowMcpToolSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "04_WorkflowMcpTool");
@@ -333,7 +333,7 @@ public sealed class WorkflowSamplesValidation(ITestOutputHelper outputHelper) :
});
}
[Fact]
[Fact(Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task WorkflowAndAgentsSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "05_WorkflowAndAgents");
@@ -385,7 +385,7 @@ public sealed class WorkflowSamplesValidation(ITestOutputHelper outputHelper) :
});
}
[Fact]
[Fact(Skip = "Azure Functions Core Tools v4 cannot auto-detect worker runtime in CI. See https://github.com/microsoft/agent-framework/issues/6402")]
public async Task ConcurrentWorkflowSampleValidationAsync()
{
string samplePath = Path.Combine(s_samplesPath, "02_ConcurrentWorkflow");
@@ -16,6 +16,7 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
private const string TestUserId = "test-user-id";
private const string CustomClaimType = "custom-claim-type";
private const string CustomClaimValue = "custom-claim-value";
private const string TestAuthenticationType = "TestAuth";
private readonly Mock<IHttpContextAccessor> _httpContextAccessorMock;
@@ -101,7 +102,7 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
public async Task GetSessionIsolationKeyAsyncExtractsDefaultClaimTypeAsync()
{
// Arrange
this.SetupHttpContextWithClaim(ClaimsIdentity.DefaultNameClaimType, TestUserId);
this.SetupHttpContextWithClaim(ClaimTypes.NameIdentifier, TestUserId);
var provider = new ClaimsIdentitySessionIsolationKeyProvider(this._httpContextAccessorMock.Object);
// Act
@@ -111,6 +112,25 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
Assert.Equal(TestUserId, result);
}
/// <summary>
/// Verify that the default claim type is the stable, unique NameIdentifier claim rather than the
/// non-unique display name claim. This guards against the session-isolation collision described in
/// the security report where two principals sharing the same name claim received the same key.
/// </summary>
[Fact]
public async Task GetSessionIsolationKeyAsyncIgnoresNameClaimByDefaultAsync()
{
// Arrange - only a display-name claim is present; the default provider must not use it.
this.SetupHttpContextWithClaim(ClaimsIdentity.DefaultNameClaimType, TestUserId);
var provider = new ClaimsIdentitySessionIsolationKeyProvider(this._httpContextAccessorMock.Object);
// Act
string? result = await provider.GetSessionIsolationKeyAsync();
// Assert
Assert.Null(result);
}
/// <summary>
/// Verify that GetSessionIsolationKeyAsync uses custom claim type when specified.
/// </summary>
@@ -191,10 +211,10 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
const string SecondValue = "second-value";
var claims = new[]
{
new Claim(ClaimsIdentity.DefaultNameClaimType, FirstValue),
new Claim(ClaimsIdentity.DefaultNameClaimType, SecondValue),
new Claim(ClaimTypes.NameIdentifier, FirstValue),
new Claim(ClaimTypes.NameIdentifier, SecondValue),
};
var identity = new ClaimsIdentity(claims);
var identity = new ClaimsIdentity(claims, TestAuthenticationType);
var principal = new ClaimsPrincipal(identity);
var httpContext = new DefaultHttpContext
@@ -219,7 +239,7 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
public async Task GetSessionIsolationKeyAsyncHandlesEmptyClaimValueAsync()
{
// Arrange
this.SetupHttpContextWithClaim(ClaimsIdentity.DefaultNameClaimType, string.Empty);
this.SetupHttpContextWithClaim(ClaimTypes.NameIdentifier, string.Empty);
var provider = new ClaimsIdentitySessionIsolationKeyProvider(this._httpContextAccessorMock.Object);
// Act
@@ -229,6 +249,66 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
Assert.Equal(string.Empty, result);
}
/// <summary>
/// Regression test for the session-isolation collision security report: two distinct authenticated
/// principals that share the same display-name claim but have different stable identifiers and tenants
/// must produce distinct isolation keys under the default options.
/// </summary>
[Fact]
public async Task GetSessionIsolationKeyAsyncDistinctForPrincipalsSharingNameClaimAsync()
{
// Arrange - both principals share the same name claim but differ by NameIdentifier and tenant.
const string CommonName = "John Doe";
var principalA = CreatePrincipal(
new Claim(ClaimsIdentity.DefaultNameClaimType, CommonName),
new Claim(ClaimTypes.NameIdentifier, "oid-user-a"),
new Claim("http://schemas.microsoft.com/identity/claims/tenantid", "tenant-a"));
var principalB = CreatePrincipal(
new Claim(ClaimsIdentity.DefaultNameClaimType, CommonName),
new Claim(ClaimTypes.NameIdentifier, "oid-user-b"),
new Claim("http://schemas.microsoft.com/identity/claims/tenantid", "tenant-b"));
var provider = new ClaimsIdentitySessionIsolationKeyProvider(this._httpContextAccessorMock.Object);
// Act
this._httpContextAccessorMock.Setup(x => x.HttpContext).Returns(new DefaultHttpContext { User = principalA });
string? principalAKey = await provider.GetSessionIsolationKeyAsync();
this._httpContextAccessorMock.Setup(x => x.HttpContext).Returns(new DefaultHttpContext { User = principalB });
string? principalBKey = await provider.GetSessionIsolationKeyAsync();
// Assert
Assert.Equal("oid-user-a", principalAKey);
Assert.Equal("oid-user-b", principalBKey);
Assert.NotEqual(principalAKey, principalBKey);
}
/// <summary>
/// Verify that GetSessionIsolationKeyAsync returns null when the request's user is not authenticated,
/// even if a claim of the configured type is present. The provider must not derive an isolation key
/// from claims on an unauthenticated identity.
/// </summary>
[Fact]
public async Task GetSessionIsolationKeyAsyncReturnsNullWhenUserNotAuthenticatedAsync()
{
// Arrange - identity has the claim but no authentication type, so IsAuthenticated is false.
var claims = new[] { new Claim(ClaimTypes.NameIdentifier, TestUserId) };
var unauthenticatedIdentity = new ClaimsIdentity(claims);
var principal = new ClaimsPrincipal(unauthenticatedIdentity);
var httpContext = new DefaultHttpContext { User = principal };
this._httpContextAccessorMock.Setup(x => x.HttpContext).Returns(httpContext);
var provider = new ClaimsIdentitySessionIsolationKeyProvider(this._httpContextAccessorMock.Object);
// Act
string? result = await provider.GetSessionIsolationKeyAsync();
// Assert
Assert.False(unauthenticatedIdentity.IsAuthenticated);
Assert.Null(result);
}
#endregion
#region Helper Methods
@@ -236,7 +316,7 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
private void SetupHttpContextWithClaim(string claimType, string claimValue)
{
var claims = new[] { new Claim(claimType, claimValue) };
var identity = new ClaimsIdentity(claims);
var identity = new ClaimsIdentity(claims, TestAuthenticationType);
var principal = new ClaimsPrincipal(identity);
var httpContext = new DefaultHttpContext
@@ -247,5 +327,8 @@ public class ClaimsIdentitySessionIsolationKeyProviderTests
this._httpContextAccessorMock.Setup(x => x.HttpContext).Returns(httpContext);
}
private static ClaimsPrincipal CreatePrincipal(params Claim[] claims)
=> new(new ClaimsIdentity(claims, TestAuthenticationType));
#endregion
}
@@ -134,6 +134,7 @@ public class ChatClientAgentOptionsTests
ClearOnChatHistoryProviderConflict = false,
WarnOnChatHistoryProviderConflict = false,
ThrowOnChatHistoryProviderConflict = false,
EnableNonApprovalRequiredFunctionBypassing = true,
};
// Act
@@ -150,6 +151,7 @@ public class ChatClientAgentOptionsTests
Assert.Equal(original.ClearOnChatHistoryProviderConflict, clone.ClearOnChatHistoryProviderConflict);
Assert.Equal(original.WarnOnChatHistoryProviderConflict, clone.WarnOnChatHistoryProviderConflict);
Assert.Equal(original.ThrowOnChatHistoryProviderConflict, clone.ThrowOnChatHistoryProviderConflict);
Assert.Equal(original.EnableNonApprovalRequiredFunctionBypassing, clone.EnableNonApprovalRequiredFunctionBypassing);
// ChatOptions should be cloned, not the same reference
Assert.NotSame(original.ChatOptions, clone.ChatOptions);
@@ -0,0 +1,574 @@
// Copyright (c) Microsoft. All rights reserved.
using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Extensions.AI;
using Moq;
namespace Microsoft.Agents.AI.UnitTests;
public class NonApprovalRequiredFunctionBypassingChatClientTests
{
#region GetResponseAsync Tests
[Fact]
public async Task GetResponseAsync_NoApprovalContent_PassesThroughUnchangedAsync()
{
// Arrange
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Hello")])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
// Act
var response = await RunWithAgentContextAsync(decorator, session);
// Assert
Assert.Single(response.Messages);
Assert.Equal("Hello", response.Messages[0].Text);
Assert.Equal(0, session.StateBag.Count);
}
[Fact]
public async Task GetResponseAsync_AllToolsRequireApproval_PassesThroughUnchangedAsync()
{
// Arrange
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));
var fcc = new FunctionCallContent("call1", "approvalTool");
var approval = new ToolApprovalRequestContent("req1", fcc);
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, [approval])])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [approvalTool] };
// Act
var response = await RunWithAgentContextAsync(decorator, session, options);
// Assert — approval request should remain
Assert.Single(response.Messages);
var contents = response.Messages[0].Contents;
Assert.Single(contents);
Assert.IsType<ToolApprovalRequestContent>(contents[0]);
Assert.Equal(0, session.StateBag.Count);
}
[Fact]
public async Task GetResponseAsync_MixedApproval_RemovesNonApprovalItemsAsync()
{
// Arrange
var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));
var fccNormal = new FunctionCallContent("call1", "normalTool");
var fccApproval = new FunctionCallContent("call2", "approvalTool");
var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([
new ChatMessage(ChatRole.Assistant, [approvalNormal, approvalRequired])
])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [normalTool, approvalTool] };
// Act
var response = await RunWithAgentContextAsync(decorator, session, options);
// Assert — only the approval-required item remains in the response
Assert.Single(response.Messages);
var contents = response.Messages[0].Contents;
Assert.Single(contents);
var remainingApproval = Assert.IsType<ToolApprovalRequestContent>(contents[0]);
Assert.Equal("req2", remainingApproval.RequestId);
}
[Fact]
public async Task GetResponseAsync_MixedApproval_StoresAutoApprovedInSessionAsync()
{
// Arrange
var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));
var fccNormal = new FunctionCallContent("call1", "normalTool");
var fccApproval = new FunctionCallContent("call2", "approvalTool");
var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([
new ChatMessage(ChatRole.Assistant, [approvalNormal, approvalRequired])
])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [normalTool, approvalTool] };
// Act
await RunWithAgentContextAsync(decorator, session, options);
// Assert — the auto-approved item should be stored in the session
Assert.True(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions));
Assert.NotNull(stored);
Assert.Single(stored!);
Assert.Equal("req1", stored![0].RequestId);
}
[Fact]
public async Task GetResponseAsync_AllNonApproval_RemovesAllApprovalsAndRemovesEmptyMessageAsync()
{
// Arrange
var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
var fccNormal = new FunctionCallContent("call1", "normalTool");
var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([
new ChatMessage(ChatRole.Assistant, [approvalNormal])
])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [normalTool] };
// Act
var response = await RunWithAgentContextAsync(decorator, session, options);
// Assert — the message should be removed since it's now empty
Assert.Empty(response.Messages);
}
[Fact]
public async Task GetResponseAsync_NextRequest_InjectsStoredAutoApprovalsAsync()
{
// Arrange
var fccNormal = new FunctionCallContent("call1", "normalTool");
var storedApproval = new ToolApprovalRequestContent("req1", fccNormal);
var session = new ChatClientAgentSession();
session.StateBag.SetValue(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey,
new List<ToolApprovalRequestContent> { storedApproval },
AgentJsonUtilities.DefaultOptions);
IEnumerable<ChatMessage>? capturedMessages = null;
var innerClient = CreateMockChatClient((messages, _, _) =>
{
capturedMessages = messages.ToList();
return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")]));
});
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var options = new ChatOptions { Tools = [AIFunctionFactory.Create(() => "result", "normalTool")] };
// Act
await RunWithAgentContextAsync(decorator, session, options);
// Assert — the inner client should receive injected messages
Assert.NotNull(capturedMessages);
var messagesList = capturedMessages!.ToList();
// Original user message + user message with approved responses.
Assert.Equal(2, messagesList.Count);
Assert.Equal(ChatRole.User, messagesList[0].Role);
// User message with the auto-approved ToolApprovalResponseContent
Assert.Equal(ChatRole.User, messagesList[1].Role);
var userContent = messagesList[1].Contents.OfType<ToolApprovalResponseContent>().ToList();
Assert.Single(userContent);
Assert.Equal("req1", userContent[0].RequestId);
Assert.True(userContent[0].Approved);
}
[Fact]
public async Task GetResponseAsync_NextRequest_ClearsStoredAfterInjectionAsync()
{
// Arrange
var fccNormal = new FunctionCallContent("call1", "normalTool");
var storedApproval = new ToolApprovalRequestContent("req1", fccNormal);
var session = new ChatClientAgentSession();
session.StateBag.SetValue(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey,
new List<ToolApprovalRequestContent> { storedApproval },
AgentJsonUtilities.DefaultOptions);
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var options = new ChatOptions { Tools = [AIFunctionFactory.Create(() => "result", "normalTool")] };
// Act
await RunWithAgentContextAsync(decorator, session, options);
// Assert — the stored data should be cleared after successful injection
Assert.False(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey, out _, AgentJsonUtilities.DefaultOptions));
}
[Fact]
public async Task GetResponseAsync_UnknownTool_TreatedAsApprovalRequiredAsync()
{
// Arrange — tool is not in ChatOptions.Tools
var fccUnknown = new FunctionCallContent("call1", "unknownTool");
var approvalUnknown = new ToolApprovalRequestContent("req1", fccUnknown);
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([
new ChatMessage(ChatRole.Assistant, [approvalUnknown])
])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [] };
// Act
var response = await RunWithAgentContextAsync(decorator, session, options);
// Assert — unknown tool should NOT be auto-approved
Assert.Single(response.Messages);
Assert.Single(response.Messages[0].Contents);
Assert.IsType<ToolApprovalRequestContent>(response.Messages[0].Contents[0]);
Assert.Equal(0, session.StateBag.Count);
}
[Fact]
public async Task GetResponseAsync_StoredRequestToolSetChanged_StillInjectsAsApprovedAsync()
{
// Arrange — tool was previously non-approval-required but is now wrapped in ApprovalRequiredAIFunction.
// The LLM still requires a complete set of responses, so we inject unconditionally.
var fccTool = new FunctionCallContent("call1", "changingTool");
var storedApproval = new ToolApprovalRequestContent("req1", fccTool);
var session = new ChatClientAgentSession();
session.StateBag.SetValue(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey,
new List<ToolApprovalRequestContent> { storedApproval },
AgentJsonUtilities.DefaultOptions);
IEnumerable<ChatMessage>? capturedMessages = null;
var innerClient = CreateMockChatClient((messages, _, _) =>
{
capturedMessages = messages.ToList();
return Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "Done")]));
});
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
// The tool is now wrapped in ApprovalRequiredAIFunction — but we still inject unconditionally
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "changingTool"));
var options = new ChatOptions { Tools = [approvalTool] };
// Act
await RunWithAgentContextAsync(decorator, session, options);
// Assert — the stored request should still be injected as approved
Assert.NotNull(capturedMessages);
var messagesList = capturedMessages!.ToList();
Assert.Equal(2, messagesList.Count);
var userContent = messagesList[1].Contents.OfType<ToolApprovalResponseContent>().ToList();
Assert.Single(userContent);
Assert.Equal("req1", userContent[0].RequestId);
Assert.True(userContent[0].Approved);
// Session should be cleared
Assert.False(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey, out _, AgentJsonUtilities.DefaultOptions));
}
#endregion
#region GetStreamingResponseAsync Tests
[Fact]
public async Task GetStreamingResponseAsync_NoApprovalContent_PassesThroughUnchangedAsync()
{
// Arrange
var innerClient = CreateMockStreamingChatClient((_, _, _) =>
ToAsyncEnumerableAsync(
new ChatResponseUpdate(ChatRole.Assistant, "Hello")));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
// Act
var updates = new List<ChatResponseUpdate>();
await RunStreamingWithAgentContextAsync(decorator, session, updates);
// Assert
Assert.Single(updates);
Assert.Equal("Hello", updates[0].Text);
Assert.Equal(0, session.StateBag.Count);
}
[Fact]
public async Task GetStreamingResponseAsync_MixedApproval_FiltersNonApprovalItemsAsync()
{
// Arrange
var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));
var fccNormal = new FunctionCallContent("call1", "normalTool");
var fccApproval = new FunctionCallContent("call2", "approvalTool");
var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);
var innerClient = CreateMockStreamingChatClient((_, _, _) =>
ToAsyncEnumerableAsync(
new ChatResponseUpdate(ChatRole.Assistant, "text"),
new ChatResponseUpdate { Contents = [approvalNormal, approvalRequired] }));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [normalTool, approvalTool] };
// Act
var updates = new List<ChatResponseUpdate>();
await RunStreamingWithAgentContextAsync(decorator, session, updates, options);
// Assert — text update + filtered approval update
Assert.Equal(2, updates.Count);
Assert.Equal("text", updates[0].Text);
// Second update should only have the approval-required item
var approvalContents = updates[1].Contents.OfType<ToolApprovalRequestContent>().ToList();
Assert.Single(approvalContents);
Assert.Equal("req2", approvalContents[0].RequestId);
}
[Fact]
public async Task GetStreamingResponseAsync_MixedApproval_StoresAutoApprovedInSessionAsync()
{
// Arrange
var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
var approvalTool = new ApprovalRequiredAIFunction(AIFunctionFactory.Create(() => "result", "approvalTool"));
var fccNormal = new FunctionCallContent("call1", "normalTool");
var fccApproval = new FunctionCallContent("call2", "approvalTool");
var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
var approvalRequired = new ToolApprovalRequestContent("req2", fccApproval);
var innerClient = CreateMockStreamingChatClient((_, _, _) =>
ToAsyncEnumerableAsync(
new ChatResponseUpdate { Contents = [approvalNormal, approvalRequired] }));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [normalTool, approvalTool] };
// Act
var updates = new List<ChatResponseUpdate>();
await RunStreamingWithAgentContextAsync(decorator, session, updates, options);
// Assert — the auto-approved item should be stored in the session
Assert.True(session.StateBag.TryGetValue<List<ToolApprovalRequestContent>>(
NonApprovalRequiredFunctionBypassingChatClient.StateBagKey, out var stored, AgentJsonUtilities.DefaultOptions));
Assert.NotNull(stored);
Assert.Single(stored!);
Assert.Equal("req1", stored![0].RequestId);
}
[Fact]
public async Task GetStreamingResponseAsync_AllNonApproval_SkipsEmptyUpdateAsync()
{
// Arrange
var normalTool = AIFunctionFactory.Create(() => "result", "normalTool");
var fccNormal = new FunctionCallContent("call1", "normalTool");
var approvalNormal = new ToolApprovalRequestContent("req1", fccNormal);
var innerClient = CreateMockStreamingChatClient((_, _, _) =>
ToAsyncEnumerableAsync(
new ChatResponseUpdate(ChatRole.Assistant, "text"),
new ChatResponseUpdate { Contents = [approvalNormal] }));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
var session = new ChatClientAgentSession();
var options = new ChatOptions { Tools = [normalTool] };
// Act
var updates = new List<ChatResponseUpdate>();
await RunStreamingWithAgentContextAsync(decorator, session, updates, options);
// Assert — the approval update should be skipped entirely
Assert.Single(updates);
Assert.Equal("text", updates[0].Text);
}
#endregion
#region Error Handling Tests
[Fact]
public async Task GetResponseAsync_NoRunContext_ThrowsInvalidOperationExceptionAsync()
{
// Arrange
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "response")])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
// Act & Assert — calling directly without agent context
await Assert.ThrowsAsync<InvalidOperationException>(
() => decorator.GetResponseAsync([new ChatMessage(ChatRole.User, "test")]));
}
[Fact]
public async Task GetResponseAsync_NoSession_ThrowsInvalidOperationExceptionAsync()
{
// Arrange
var innerClient = CreateMockChatClient((_, _, _) =>
Task.FromResult(new ChatResponse([new ChatMessage(ChatRole.Assistant, "response")])));
var decorator = new NonApprovalRequiredFunctionBypassingChatClient(innerClient);
// Act & Assert — run with null session
await Assert.ThrowsAsync<InvalidOperationException>(
() => RunWithAgentContextAsync(decorator, session: null!));
}
#endregion
#region Builder Extension Tests
[Fact]
public void UseNonApprovalRequiredFunctionBypassing_AddsDecoratorToPipeline()
{
// Arrange
var innerClient = new Mock<IChatClient>().Object;
// Act
var pipeline = innerClient.AsBuilder()
.UseNonApprovalRequiredFunctionBypassing()
.Build();
// Assert
Assert.NotNull(pipeline.GetService<NonApprovalRequiredFunctionBypassingChatClient>());
}
[Fact]
public void WithDefaultAgentMiddleware_EnableNonApprovalRequiredFunctionBypassing_InjectsDecorator()
{
// Arrange
var innerClient = new Mock<IChatClient>().Object;
var options = new ChatClientAgentOptions { EnableNonApprovalRequiredFunctionBypassing = true };
// Act
var pipeline = innerClient.WithDefaultAgentMiddleware(options);
// Assert
Assert.NotNull(pipeline.GetService<NonApprovalRequiredFunctionBypassingChatClient>());
}
[Fact]
public void WithDefaultAgentMiddleware_EnableNonApprovalRequiredFunctionBypassingFalse_DoesNotInjectDecorator()
{
// Arrange
var innerClient = new Mock<IChatClient>().Object;
var options = new ChatClientAgentOptions { EnableNonApprovalRequiredFunctionBypassing = false };
// Act
var pipeline = innerClient.WithDefaultAgentMiddleware(options);
// Assert
Assert.Null(pipeline.GetService<NonApprovalRequiredFunctionBypassingChatClient>());
}
#endregion
#region Helpers
private static async Task<ChatResponse> RunWithAgentContextAsync(
NonApprovalRequiredFunctionBypassingChatClient decorator,
AgentSession? session,
ChatOptions? options = null)
{
ChatResponse? capturedResponse = null;
var agent = new TestAIAgent
{
RunAsyncFunc = async (messages, agentSession, agentOptions, ct) =>
{
capturedResponse = await decorator.GetResponseAsync(messages, options, ct);
return new AgentResponse(capturedResponse);
}
};
await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], session);
return capturedResponse!;
}
private static Task<ChatResponse> RunWithAgentContextAsync(
NonApprovalRequiredFunctionBypassingChatClient decorator,
AgentSession session)
=> RunWithAgentContextAsync(decorator, session, options: null);
private static async Task RunStreamingWithAgentContextAsync(
NonApprovalRequiredFunctionBypassingChatClient decorator,
AgentSession session,
List<ChatResponseUpdate> updates,
ChatOptions? options = null)
{
var agent = new TestAIAgent
{
RunAsyncFunc = async (messages, agentSession, agentOptions, ct) =>
{
await foreach (var update in decorator.GetStreamingResponseAsync(messages, options, ct))
{
updates.Add(update);
}
return new AgentResponse([new ChatMessage(ChatRole.Assistant, "done")]);
}
};
await agent.RunAsync([new ChatMessage(ChatRole.User, "Hello")], session);
}
private static IChatClient CreateMockChatClient(
Func<IEnumerable<ChatMessage>, ChatOptions?, CancellationToken, Task<ChatResponse>> onGetResponse)
{
var mock = new Mock<IChatClient>();
mock.Setup(c => c.GetResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions?>(),
It.IsAny<CancellationToken>()))
.Returns((IEnumerable<ChatMessage> m, ChatOptions? o, CancellationToken ct) => onGetResponse(m, o, ct));
return mock.Object;
}
private static IChatClient CreateMockStreamingChatClient(
Func<IEnumerable<ChatMessage>, ChatOptions?, CancellationToken, IAsyncEnumerable<ChatResponseUpdate>> onGetStreamingResponse)
{
var mock = new Mock<IChatClient>();
mock.Setup(c => c.GetStreamingResponseAsync(
It.IsAny<IEnumerable<ChatMessage>>(),
It.IsAny<ChatOptions?>(),
It.IsAny<CancellationToken>()))
.Returns((IEnumerable<ChatMessage> m, ChatOptions? o, CancellationToken ct) => onGetStreamingResponse(m, o, ct));
return mock.Object;
}
private static async IAsyncEnumerable<ChatResponseUpdate> ToAsyncEnumerableAsync(params ChatResponseUpdate[] updates)
{
foreach (var update in updates)
{
yield return update;
}
await Task.CompletedTask;
}
#endregion
}
@@ -59,15 +59,15 @@ public class ToolApprovalAgentBuilderExtensionsTests
/// Verify that UseToolApproval with custom JsonSerializerOptions works correctly.
/// </summary>
[Fact]
public void UseToolApproval_WithCustomJsonSerializerOptions_ReturnsToolApprovalAgent()
public void UseToolApproval_WithCustomOptions_ReturnsToolApprovalAgent()
{
// Arrange
var mockAgent = new Mock<AIAgent>();
var builder = new AIAgentBuilder(mockAgent.Object);
var options = new JsonSerializerOptions();
var options = new ToolApprovalAgentOptions { JsonSerializerOptions = new JsonSerializerOptions() };
// Act
var result = builder.UseToolApproval(jsonSerializerOptions: options).Build();
var result = builder.UseToolApproval(options: options).Build();
// Assert
Assert.IsType<ToolApprovalAgent>(result);
@@ -47,14 +47,14 @@ public class ToolApprovalAgentTests
}
/// <summary>
/// Verify that constructor accepts custom JsonSerializerOptions.
/// Verify that constructor accepts custom options.
/// </summary>
[Fact]
public void Constructor_CustomJsonSerializerOptions_CreatesInstanceAsync()
public void Constructor_CustomOptions_CreatesInstance()
{
// Arrange
var innerAgent = new Mock<AIAgent>().Object;
var options = new JsonSerializerOptions();
var options = new ToolApprovalAgentOptions { JsonSerializerOptions = new JsonSerializerOptions() };
// Act
var agent = new ToolApprovalAgent(innerAgent, options);
@@ -1535,4 +1535,311 @@ public class ToolApprovalAgentTests
}
#endregion
#region Auto-Approval Rules (Heuristics)
/// <summary>
/// Verify that an auto-approval rule can approve a function call that would otherwise need user approval.
/// </summary>
[Fact]
public async Task RunAsync_AutoApprovalRule_ApprovesMatchingToolAsync()
{
// Arrange
var session = new ChatClientAgentSession();
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadTool"));
// Inner agent: first call returns approval request, second returns final response.
var callCount = 0;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(() =>
{
callCount++;
if (callCount == 1)
{
return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]);
}
return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]);
});
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(fcc.Name == "ReadTool")]
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Act
var response = await agent.RunAsync(
[new ChatMessage(ChatRole.User, "Hi")],
session);
// Assert — the approval request was auto-approved, inner agent called twice
Assert.Equal(2, callCount);
Assert.Equal("Done", response.Text);
}
/// <summary>
/// Verify that when auto-approval rule does not match, request is surfaced to the caller.
/// </summary>
[Fact]
public async Task RunAsync_AutoApprovalRule_DoesNotMatchSurfacesToCallerAsync()
{
// Arrange
var session = new ChatClientAgentSession();
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "DangerousTool"));
var innerAgent = CreateMockAgent(new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]));
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(fcc.Name == "ReadTool")] // Only approves ReadTool
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Act
var response = await agent.RunAsync(
[new ChatMessage(ChatRole.User, "Hi")],
session);
// Assert — request surfaced to caller since heuristic doesn't match
var requests = response.Messages.SelectMany(m => m.Contents).OfType<ToolApprovalRequestContent>().ToList();
Assert.Single(requests);
Assert.Equal("DangerousTool", ((FunctionCallContent)requests[0].ToolCall).Name);
}
/// <summary>
/// Verify that multiple auto-approval rules are evaluated in order; first match wins.
/// </summary>
[Fact]
public async Task RunAsync_MultipleAutoApprovalRules_FirstMatchWinsAsync()
{
// Arrange
var session = new ChatClientAgentSession();
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "SpecialTool"));
var callCount = 0;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(() =>
{
callCount++;
if (callCount == 1)
{
return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]);
}
return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]);
});
var rule1Called = false;
var rule2Called = false;
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules =
[
fcc => { rule1Called = true; return new ValueTask<bool>(fcc.Name == "SpecialTool"); },
fcc => { rule2Called = true; return new ValueTask<bool>(true); } // Should not be reached
]
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Act
await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session);
// Assert — first rule matched, second was never called
Assert.True(rule1Called);
Assert.False(rule2Called);
}
/// <summary>
/// Verify that standing rules are evaluated before auto-approval rules.
/// </summary>
[Fact]
public async Task RunAsync_StandingRuleTakesPrecedenceOverAutoApprovalRuleAsync()
{
// Arrange
var session = new ChatClientAgentSession();
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "MyTool"));
var callCount = 0;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.ReturnsAsync(() =>
{
callCount++;
if (callCount <= 2)
{
return new AgentResponse([new ChatMessage(ChatRole.Assistant, [approvalRequest])]);
}
return new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]);
});
var heuristicCalled = false;
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => { heuristicCalled = true; return new ValueTask<bool>(true); }]
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Call 1: heuristic should be called (no standing rule yet)
var response1 = await agent.RunAsync([new ChatMessage(ChatRole.User, "Hi")], session);
Assert.True(heuristicCalled);
Assert.Equal("Done", response1.Text);
// Now establish a standing rule by sending AlwaysApprove
heuristicCalled = false;
callCount = 0;
var alwaysApprove = new AlwaysApproveToolApprovalResponseContent(
approvalRequest.CreateResponse(approved: true),
alwaysApproveTool: true,
alwaysApproveToolWithArguments: false);
// Call 2: standing rule should match first, heuristic should NOT be called
var response2 = await agent.RunAsync(
[new ChatMessage(ChatRole.User, [alwaysApprove])],
session);
Assert.False(heuristicCalled);
Assert.Equal("Done", response2.Text);
}
/// <summary>
/// Verify that when a batch contains a mix of heuristic-approved and standing-rule-approved
/// requests, the collected approval responses preserve the original request order rather than
/// being grouped by approval kind.
/// </summary>
[Fact]
public async Task RunAsync_MixedAutoApprovals_PreserveOriginalOrderAsync()
{
// Arrange
var session = new ChatClientAgentSession();
// Batch ordering: first request is approved by a heuristic, second by a standing rule.
var heuristicRequest = new ToolApprovalRequestContent("reqA", new FunctionCallContent("callA", "HeuristicTool"));
var standingRequest = new ToolApprovalRequestContent("reqB", new FunctionCallContent("callB", "StandingTool"));
var batchResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, [heuristicRequest, standingRequest])]);
var finalResponse = new AgentResponse([new ChatMessage(ChatRole.Assistant, "Done")]);
var callCount = 0;
List<ChatMessage>? secondCallMessages = null;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<Task<AgentResponse>>("RunCoreAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.Callback<IEnumerable<ChatMessage>, AgentSession?, AgentRunOptions?, CancellationToken>((msgs, _, _, _) =>
{
callCount++;
if (callCount == 2)
{
secondCallMessages = msgs.ToList();
}
})
.ReturnsAsync(() => callCount == 1 ? batchResponse : finalResponse);
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(fcc.Name == "HeuristicTool")]
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Establish a standing rule for "StandingTool" via an AlwaysApprove response in the same call.
var alwaysApprove = standingRequest.CreateAlwaysApproveToolResponse("User said always");
// Act — both requests auto-approve (heuristic + standing rule), so the inner agent is re-invoked.
var response = await agent.RunAsync(
[new ChatMessage(ChatRole.User, [alwaysApprove])],
session);
// Assert — inner agent re-called and final response returned.
Assert.Equal(2, callCount);
Assert.Equal("Done", response.Text);
// The injected approval responses must preserve the original request order: reqA before reqB,
// even though reqA was approved by a heuristic and reqB by a standing rule.
Assert.NotNull(secondCallMessages);
var injected = secondCallMessages!
.SelectMany(m => m.Contents)
.OfType<ToolApprovalResponseContent>()
.Where(r => r.RequestId is "reqA" or "reqB")
.ToList();
Assert.Equal(2, injected.Count);
Assert.Equal("reqA", injected[0].RequestId);
Assert.Equal("reqB", injected[1].RequestId);
Assert.All(injected, r => Assert.True(r.Approved));
}
/// <summary>
/// Verify that auto-approval rules work in the streaming path.
/// </summary>
[Fact]
public async Task RunStreamingAsync_AutoApprovalRule_ApprovesMatchingToolAsync()
{
// Arrange
var session = new ChatClientAgentSession();
var approvalRequest = new ToolApprovalRequestContent("req1", new FunctionCallContent("call1", "ReadTool"));
var callCount = 0;
var innerAgent = new Mock<AIAgent>();
innerAgent
.Protected()
.Setup<IAsyncEnumerable<AgentResponseUpdate>>("RunCoreStreamingAsync",
ItExpr.IsAny<IEnumerable<ChatMessage>>(),
ItExpr.IsAny<AgentSession?>(),
ItExpr.IsAny<AgentRunOptions?>(),
ItExpr.IsAny<CancellationToken>())
.Returns(() =>
{
callCount++;
if (callCount == 1)
{
return ToAsyncEnumerableAsync([new AgentResponseUpdate(ChatRole.Assistant, [approvalRequest])]);
}
return ToAsyncEnumerableAsync([new AgentResponseUpdate(ChatRole.Assistant, "Done")]);
});
var options = new ToolApprovalAgentOptions
{
AutoApprovalRules = [fcc => new ValueTask<bool>(fcc.Name == "ReadTool")]
};
var agent = new ToolApprovalAgent(innerAgent.Object, options);
// Act
var updates = new List<AgentResponseUpdate>();
await foreach (var update in agent.RunStreamingAsync([new ChatMessage(ChatRole.User, "Hi")], session))
{
updates.Add(update);
}
// Assert — the approval request was auto-approved, inner agent streamed twice
Assert.Equal(2, callCount);
Assert.Single(updates);
Assert.Equal("Done", updates[0].Text);
}
#endregion
}
@@ -142,6 +142,95 @@ public sealed class ForeachExecutorTest(ITestOutputHelper output) : WorkflowActi
indexName: "CurrentIndex");
}
[Fact]
public async Task ForeachTakeNextWithMultiFieldRecordAsync()
{
// Arrange
const string CurrentValueName = "CurrentValue";
this.SetVariableState(CurrentValueName);
TableDataValue tableValue = DataValue.TableFromRecords(
DataValue.RecordFromFields(
new KeyValuePair<string, DataValue>("name", new StringDataValue("Alice")),
new KeyValuePair<string, DataValue>("role", new StringDataValue("Engineer"))));
Foreach model = this.CreateModel(
displayName: nameof(ForeachTakeNextWithMultiFieldRecordAsync),
items: ValueExpression.Literal(tableValue),
valueName: CurrentValueName,
indexName: null);
ForeachExecutor action = new(model, this.State);
// Act
await this.ExecuteAsync(action, ForeachExecutor.Steps.Next(action.Id), action.TakeNextAsync);
// Assert
RecordValue currentValue = Assert.IsType<RecordValue>(this.State.Get(CurrentValueName), exactMatch: false);
Assert.Equal("Alice", currentValue.GetField("name").ToObject());
Assert.Equal("Engineer", currentValue.GetField("role").ToObject());
}
/// <summary>
/// Power Fx wraps scalar array literals such as <c>=[1, 2, 3]</c> as <c>Table({Value: 1}, ...)</c>;
/// the loop value must expose the bare scalar, not the single-column wrapper record.
/// </summary>
[Fact]
public async Task ForeachTakeNextWithSingleColumnValueRecordAsync()
{
// Arrange
const string CurrentValueName = "CurrentValue";
this.SetVariableState(CurrentValueName);
TableDataValue tableValue = DataValue.TableFromRecords(
DataValue.RecordFromFields(new KeyValuePair<string, DataValue>("Value", new NumberDataValue(1))),
DataValue.RecordFromFields(new KeyValuePair<string, DataValue>("Value", new NumberDataValue(2))),
DataValue.RecordFromFields(new KeyValuePair<string, DataValue>("Value", new NumberDataValue(3))));
Foreach model = this.CreateModel(
displayName: nameof(ForeachTakeNextWithSingleColumnValueRecordAsync),
items: ValueExpression.Literal(tableValue),
valueName: CurrentValueName,
indexName: null);
ForeachExecutor action = new(model, this.State);
// Act
await this.ExecuteAsync(action, ForeachExecutor.Steps.Next(action.Id), action.TakeNextAsync);
// Assert
FormulaValue currentValue = this.State.Get(CurrentValueName);
Assert.IsNotType<RecordValue>(currentValue, exactMatch: false);
Assert.Equal(1m, currentValue.ToObject());
}
/// <summary>
/// Single-field records whose only field is NOT named <c>Value</c> are not Power Fx auto-wraps;
/// they are preserved as records so the field name remains accessible inside the loop body.
/// </summary>
[Fact]
public async Task ForeachTakeNextWithSingleFieldNonValueRecordAsync()
{
// Arrange
const string CurrentValueName = "CurrentValue";
this.SetVariableState(CurrentValueName);
TableDataValue tableValue = DataValue.TableFromRecords(
DataValue.RecordFromFields(new KeyValuePair<string, DataValue>("name", new StringDataValue("Alice"))));
Foreach model = this.CreateModel(
displayName: nameof(ForeachTakeNextWithSingleFieldNonValueRecordAsync),
items: ValueExpression.Literal(tableValue),
valueName: CurrentValueName,
indexName: null);
ForeachExecutor action = new(model, this.State);
// Act
await this.ExecuteAsync(action, ForeachExecutor.Steps.Next(action.Id), action.TakeNextAsync);
// Assert
RecordValue currentValue = Assert.IsType<RecordValue>(this.State.Get(CurrentValueName), exactMatch: false);
Assert.Equal("Alice", currentValue.GetField("name").ToObject());
}
[Fact]
public async Task ForeachTakeLastAsync()
{
@@ -9,7 +9,7 @@ from typing import Any, cast
from ag_ui.core import BaseEvent
from agent_framework import SupportsAgentRun
from ._agent_run import run_agent_stream
from ._agent_run import PendingApprovalEntry, run_agent_stream
class AgentConfig:
@@ -107,7 +107,7 @@ class AgentFrameworkAgent:
# Populated when approval requests are emitted; consumed when responses arrive.
# Prevents bypass, function name spoofing, and replay attacks.
# Bounded to prevent unbounded growth from abandoned approval requests.
self._pending_approvals: OrderedDict[str, str] = OrderedDict()
self._pending_approvals: OrderedDict[str, PendingApprovalEntry] = OrderedDict()
self._pending_approvals_max_size: int = 10_000
async def run(
@@ -8,7 +8,7 @@ import json
import logging
import uuid
from collections.abc import AsyncIterable, Awaitable
from typing import TYPE_CHECKING, Any, cast
from typing import TYPE_CHECKING, Any, TypedDict, cast
from ag_ui.core import (
BaseEvent,
@@ -56,6 +56,7 @@ from ._run_common import (
_stringify_tool_result, # type: ignore
)
from ._utils import (
canonical_function_arguments,
convert_agui_tools_to_agent_framework,
generate_event_id,
get_conversation_id_from_update,
@@ -407,7 +408,33 @@ def _make_approval_tool_result_events(resolved_approval_results: list[Content])
return events
def _evict_oldest_approvals(registry: dict[str, str], max_size: int = 10_000) -> None:
class _PendingApproval(TypedDict):
"""Pending approval details for a requested function call."""
name: str
arguments: str | None
PendingApprovalEntry = _PendingApproval | str
def _make_pending_approval_entry(name: str, arguments: str | None) -> _PendingApproval:
return {"name": name, "arguments": arguments}
def _pending_approval_name(entry: PendingApprovalEntry) -> str | None:
if isinstance(entry, str):
return entry
return entry["name"]
def _pending_approval_arguments(entry: PendingApprovalEntry) -> str | None:
if isinstance(entry, str):
return None
return entry["arguments"]
def _evict_oldest_approvals(registry: dict[str, PendingApprovalEntry], max_size: int = 10_000) -> None:
"""Evict the oldest entries from the pending-approvals registry (LRU).
Only effective when *registry* is an ``OrderedDict``; plain dicts are
@@ -427,7 +454,7 @@ async def _resolve_approval_responses(
tools: list[Any],
agent: SupportsAgentRun,
run_kwargs: dict[str, Any],
pending_approvals: dict[str, str] | None = None,
pending_approvals: dict[str, PendingApprovalEntry] | None = None,
thread_id: str = "",
) -> list[Content]:
"""Execute approved function calls and replace approval content with results.
@@ -480,7 +507,8 @@ async def _resolve_approval_responses(
invalid_ids.add(resp_id)
continue
pending_name = pending_approvals[registry_key]
pending_entry = pending_approvals[registry_key]
pending_name = _pending_approval_name(pending_entry)
if resp_name != pending_name:
logger.warning(
"Rejected approval response id=%s: function name mismatch (response=%s, pending=%s)",
@@ -491,6 +519,16 @@ async def _resolve_approval_responses(
invalid_ids.add(resp_id)
continue
pending_arguments = _pending_approval_arguments(pending_entry)
response_arguments = canonical_function_arguments(resp.function_call)
if pending_arguments is not None and response_arguments != pending_arguments:
logger.warning(
"Rejected approval response id=%s: function arguments mismatch",
resp_id,
)
invalid_ids.add(resp_id)
continue
# Valid — consume entry to prevent replay
del pending_approvals[registry_key]
if resp.approved:
@@ -714,7 +752,7 @@ async def run_agent_stream(
input_data: dict[str, Any],
agent: SupportsAgentRun,
config: AgentConfig,
pending_approvals: dict[str, str] | None = None,
pending_approvals: dict[str, PendingApprovalEntry] | None = None,
) -> AsyncGenerator[BaseEvent]:
"""Run agent and yield AG-UI events.
@@ -917,7 +955,10 @@ async def run_agent_stream(
# Register pending approval requests so we can validate responses later
if content_type == "function_approval_request" and pending_approvals is not None:
if content.id and content.function_call and content.function_call.name:
pending_approvals[f"{thread_id}:{content.id}"] = content.function_call.name
pending_approvals[f"{thread_id}:{content.id}"] = _make_pending_approval_entry(
content.function_call.name,
canonical_function_arguments(content.function_call),
)
# Evict oldest entries if the registry exceeds a safe bound (LRU)
_evict_oldest_approvals(pending_approvals, max_size=10_000)
else:
@@ -56,6 +56,22 @@ def safe_json_parse(value: Any) -> dict[str, Any] | None:
return None
def canonical_function_arguments(function_call: Any) -> str | None:
"""Return a stable representation of function-call arguments."""
if function_call is None:
return None
try:
parsed_arguments = function_call.parse_arguments()
except Exception:
parsed_arguments = getattr(function_call, "arguments", None)
if parsed_arguments is None:
parsed_arguments = {}
return json.dumps(make_json_safe(parsed_arguments), sort_keys=True, separators=(",", ":"))
def get_role_value(message: Any) -> str:
"""Extract role string from a message object.
@@ -35,7 +35,7 @@ from ._run_common import (
_extract_resume_payload,
_normalize_resume_interrupts,
)
from ._utils import generate_event_id, make_json_safe
from ._utils import canonical_function_arguments, generate_event_id, make_json_safe
logger = logging.getLogger(__name__)
@@ -324,6 +324,29 @@ def _coerce_response_for_request(request_event: Any, value: Any) -> Any | None:
return candidate
def _approval_response_matches_request(request_id: str, request_event: Any, response: Any) -> bool:
"""Check whether an approval response matches the pending approval request."""
request_data = getattr(request_event, "data", None)
if not isinstance(request_data, Content) or request_data.type != "function_approval_request":
return True
if not isinstance(response, Content) or response.type != "function_approval_response":
return False
if str(getattr(response, "id", "")) != request_id:
return False
request_call = getattr(request_data, "function_call", None)
response_call = getattr(response, "function_call", None)
if request_call is None or response_call is None:
return False
if getattr(response_call, "name", None) != getattr(request_call, "name", None):
return False
return canonical_function_arguments(response_call) == canonical_function_arguments(request_call)
def _single_pending_response_from_value(pending_events: dict[str, Any], value: Any) -> dict[str, Any]:
"""Map a scalar resume payload to the single pending request (if unambiguous)."""
if value is None or len(pending_events) != 1:
@@ -343,6 +366,13 @@ def _single_pending_response_from_value(pending_events: dict[str, Any], value: A
)
return {}
if not _approval_response_matches_request(str(request_id), request_event, coerced_value):
logger.info(
"Ignoring pending request response for request_id=%s: approval response does not match pending request",
request_id,
)
return {}
return {str(request_id): coerced_value}
@@ -372,6 +402,12 @@ def _coerce_responses_for_pending_requests(
_response_type_name(request_event),
)
continue
if not _approval_response_matches_request(request_key, request_event, coerced_value):
logger.info(
"Ignoring resume response for request_id=%s: approval response does not match pending request",
request_key,
)
continue
normalized[request_key] = coerced_value
return normalized
@@ -1407,6 +1407,92 @@ async def test_fabricated_rejection_without_pending_approval_is_blocked(streamin
assert False, "Fabricated rejection response leaked as function_result into LLM messages"
async def test_approval_argument_mismatch_is_blocked(streaming_chat_client_stub):
"""An approval response must not execute changed arguments for the pending call."""
from agent_framework import tool
from agent_framework.ag_ui import AgentFrameworkAgent
executed_args: list[dict[str, Any]] = []
@tool(
name="update_record",
description="Update a record",
approval_mode="always_require",
)
def update_record(record_id: str, value: str) -> str:
executed_args.append({"record_id": record_id, "value": value})
return f"updated {record_id} to {value}"
async def stream_fn_approval(
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(
contents=[
Content.from_function_call(
name="update_record",
call_id="call_update_001",
arguments={"record_id": "alpha", "value": "approved"},
)
]
)
wrapper = AgentFrameworkAgent(
agent=Agent(
client=streaming_chat_client_stub(stream_fn_approval),
name="test_agent",
instructions="Test",
tools=[update_record],
)
)
thread_id = "thread-argument-mismatch-test"
events1: list[Any] = []
async for event in wrapper.run({"thread_id": thread_id, "messages": [{"role": "user", "content": "update"}]}):
events1.append(event)
assert any("call_update_001" in k for k in wrapper._pending_approvals)
async def stream_fn_post(
messages: MutableSequence[Message], options: ChatOptions, **kwargs: Any
) -> AsyncIterator[ChatResponseUpdate]:
yield ChatResponseUpdate(contents=[Content.from_text(text="Done")])
wrapper.agent = Agent(
client=streaming_chat_client_stub(stream_fn_post),
name="test_agent",
instructions="Test",
tools=[update_record],
)
turn2_input: dict[str, Any] = {
"thread_id": thread_id,
"messages": [
{
"role": "user",
"content": "approve",
"function_approvals": [
{
"id": "call_update_001",
"call_id": "call_update_001",
"name": "update_record",
"approved": True,
"arguments": {"record_id": "beta", "value": "changed"},
}
],
},
],
}
events2: list[Any] = []
async for event in wrapper.run(turn2_input):
events2.append(event)
assert executed_args == []
assert any("call_update_001" in k for k in wrapper._pending_approvals), (
"Pending approval should be preserved after argument mismatch for legitimate retry"
)
async def test_state_update_end_to_end_via_real_tool_invocation(streaming_chat_client_stub):
"""End-to-end coverage for issue #3167: a real ``@tool`` returning ``state_update`` must
emit a deterministic STATE_SNAPSHOT through the full pipeline.
@@ -1352,6 +1352,70 @@ async def test_workflow_run_approval_via_messages_approved() -> None:
assert not resumed_finished.get("interrupt")
async def test_workflow_run_approval_argument_mismatch_keeps_interrupt_pending() -> None:
"""Workflow approval responses must not resume with changed function arguments."""
handled_responses: list[dict[str, Any]] = []
class ApprovalExecutor(Executor):
def __init__(self) -> None:
super().__init__(id="approval_executor")
@handler
async def start(self, message: Any, ctx: WorkflowContext) -> None:
del message
function_call = Content.from_function_call(
call_id="refund-call",
name="submit_refund",
arguments={"order_id": "12345", "amount": "$89.99"},
)
approval_request = Content.from_function_approval_request(id="approval-1", function_call=function_call)
await ctx.request_info(approval_request, Content, request_id="approval-1")
@response_handler
async def handle_approval(self, original_request: Content, response: Content, ctx: WorkflowContext) -> None:
del original_request
if response.function_call is not None:
handled_responses.append(response.function_call.parse_arguments() or {})
await ctx.yield_output("handled")
workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build()
first_events = [
event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow)
]
first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0].model_dump()
interrupt_payload = cast(list[dict[str, Any]], first_finished.get("interrupt"))
assert isinstance(interrupt_payload, list) and len(interrupt_payload) == 1
resumed_events = [
event
async for event in run_workflow_stream(
{
"messages": [
{
"role": "user",
"content": "",
"function_approvals": [
{
"approved": True,
"id": "approval-1",
"call_id": "refund-call",
"name": "submit_refund",
"arguments": {"order_id": "99999", "amount": "$1000.00"},
}
],
}
],
},
workflow,
)
]
assert handled_responses == []
resumed_finished = [event for event in resumed_events if event.type == "RUN_FINISHED"][0].model_dump()
assert resumed_finished.get("interrupt")
async def test_workflow_run_approval_via_messages_denied() -> None:
"""Denied approval response sent via messages (function_approvals) should satisfy the pending request."""
+13
View File
@@ -76,6 +76,19 @@ agent_framework/
- **`SkillScriptRunner`** - Protocol for file-based script execution. Any callable matching `(skill, script, args) -> Any` satisfies it. Code-defined scripts do not use a runner.
- **`SkillsProvider`** - Context provider (extends `ContextProvider`) that discovers file-based skills from `SKILL.md` files and/or accepts code-defined `Skill` instances. Follows progressive disclosure: advertise → load → read resources / run scripts.
### Model Context Protocol (`_mcp.py`)
- **`MCPTool`** - Base wrapper that owns the MCP `ClientSession` and exposes the remote server's tools as `FunctionTool`s.
- **`MCPStdioTool`** / **`MCPStreamableHTTPTool`** / **`MCPWebsocketTool`** - Transport-specific subclasses.
- **`MCPTaskOptions`** (experimental, `MCP_LONG_RUNNING_TASKS` feature, **frozen**) - Per-tool-instance options controlling the SEP-2663 long-running task lifecycle. When the server advertises a tool with `execution.taskSupport == "required"`, `MCPTool.call_tool` transparently routes through `call_tool_as_task`, which sends an augmented `tools/call`, polls `tasks/get` until terminal, and reinterprets `tasks/result` as a normal `CallToolResult`. Instances are immutable; replace via `MCPTool.task_options = MCPTaskOptions(...)`. Fields:
- `default_ttl: timedelta | None` — forwarded to the server as `params.task.ttl` (milliseconds). When `None`, the server's default applies.
- `cancel_remote_task_on_local_cancellation: bool = True` — only gates the `CancelledError` path. Abandonment paths (see below) always cancel.
- `max_task_wait: timedelta | None` — client-side deadline for the whole post-create lifecycle (poll + result fetch). When exceeded, raises `ToolExecutionException` and fires a best-effort `tasks/cancel`. `None` (default) means no client-side bound. Bounds sleeps, sends, AND reconnects via `asyncio.wait_for`.
- **Permissive fallback**: servers that ignore the augmentation (return `CallToolResult` directly) or reject the unknown `task` field with `METHOD_NOT_FOUND` / `INVALID_PARAMS` fall back to the plain `session.call_tool(...)` path so legacy servers keep working. An unparseable success response (server accepted the augmented call but returned a payload that is neither `CreateTaskResult` nor `CallToolResult`) **does not** fall back — it raises `ToolExecutionException` to avoid double-executing a side-effecting tool.
- **Submit-vs-track reconnect policy**: a dropped connection before a `task_id` is known raises `ToolExecutionException("connection lost; task state unknown")` without re-issuing the augmented `tools/call`, so a server that accepted the request but lost the response cannot be made to start the same operation twice; once a `task_id` exists, `tasks/get` / `tasks/result` reconnect once and retry against the same id (a shared `_send_with_one_reconnect` helper).
- **Cancel-on-abandonment vs terminal failure**: any path where the remote task may still be running (max-wait exceeded, hard `McpError` in poll, malformed `tasks/get`, second connection loss in poll/fetch, reconnect failure) fires best-effort `tasks/cancel` before raising. Terminal failures (`failed`/`cancelled`/`input_required` server-side, `completed+isError`, malformed `tasks/result` after server completed) do **not** cancel — the server is already done. `_MCPTaskAbandoned` is the private marker distinguishing the two.
- **Transient poll retry**: a slow `tasks/get` that surfaces as `McpError(code=408 REQUEST_TIMEOUT)` is retried (bounded by `max_task_wait`). All other non-connection `McpError`s during poll are treated as abandonment. `tasks/result` does not get transient retry — the server has already completed, so a slow payload fetch is anomalous.
### File Access Harness (`_harness/_file_access.py`)
- **`AgentFileStore`** - Abstract async store backing the file-access harness. Implementations expose `write_file`, `read_file`, `delete_file`, `list_files`, `file_exists`, `search_files`, and `create_directory` over forward-slash relative paths.
@@ -124,7 +124,7 @@ from ._harness._todo import (
TodoSessionStore,
TodoStore,
)
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool
from ._mcp import MCPStdioTool, MCPStreamableHTTPTool, MCPTaskOptions, MCPWebsocketTool
from ._middleware import (
AgentContext,
AgentMiddleware,
@@ -444,12 +444,13 @@ __all__ = [
"InlineSkillResource",
"InlineSkillScript",
"LocalEvaluator",
"MCPStdioTool",
"MCPStreamableHTTPTool",
"MCPWebsocketTool",
"MCPSkill",
"MCPSkillResource",
"MCPSkillsSource",
"MCPStdioTool",
"MCPStreamableHTTPTool",
"MCPTaskOptions",
"MCPWebsocketTool",
"MemoryContextProvider",
"MemoryFileStore",
"MemoryIndexEntry",
+80 -29
View File
@@ -92,12 +92,16 @@ OptionsCoT = TypeVar(
def _merge_options(base: dict[str, Any], override: dict[str, Any]) -> dict[str, Any]:
"""Merge two options dicts, with override values taking precedence.
``None`` is treated as "unset": ``None`` overrides are skipped so they don't clobber a base
value, and the merged result is stripped of any remaining ``None`` values in a final pass so
unset options are never forwarded (e.g. an unset ``store`` is left for the service to default).
Args:
base: The base options dict.
override: The override options dict (values take precedence).
Returns:
A new merged options dict.
A new merged options dict containing no ``None`` values.
"""
result = dict(base)
@@ -123,7 +127,7 @@ def _merge_options(base: dict[str, Any], override: dict[str, Any]) -> dict[str,
result["instructions"] = f"{result['instructions']}\n{value}"
else:
result[key] = value
return result
return {key: value for key, value in result.items() if value is not None}
def _sanitize_agent_name(agent_name: str | None) -> str | None:
@@ -460,6 +464,9 @@ class BaseAgent(SerializationMixin):
if provider_session is None and self.context_providers:
provider_session = AgentSession()
# When per-service-call persistence is enabled, the per-service-call middleware owns
# HistoryProvider persistence (in both the local and service-managed cases), so skip
# them on the once-per-run path to avoid double persistence.
per_service_call_history_required = self.require_per_service_call_history_persistence and any(
isinstance(provider, HistoryProvider) for provider in self.context_providers
)
@@ -686,11 +693,16 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
description: A brief description of the agent's purpose.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
require_per_service_call_history_persistence: When True, history providers are invoked
around each model call instead of once per ``run()`` when the service
is not already storing history. If service-side storage is active for
the run, the agent skips local history providers and relies on the
service-managed conversation instead.
require_per_service_call_history_persistence: When True (and a HistoryProvider is
present), the provider always persists history via per-service-call middleware,
regardless of whether the client stores history server-side. If the client does
not store history, the middleware also loads providers around each model call and
drives the function loop with a local conversation; if it does, loading is skipped
(the service-managed conversation is the source of truth) and the middleware only
persists. A warning is logged for providers with ``load_messages=True`` when
loading is skipped because service-side storage is active. When no HistoryProvider
is present, this flag has no effect (no middleware is installed and nothing is
persisted).
default_options: A TypedDict containing chat options. When using a typed agent like
``Agent[OpenAIChatOptions]``, this enables IDE autocomplete for
provider-specific options including temperature, max_tokens, model,
@@ -791,22 +803,20 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
self,
*,
session: AgentSession | None,
options: Mapping[str, Any] | None,
conversation_id: str | None,
service_stores_history: bool,
) -> list[HistoryProvider]:
history_providers = self._get_history_providers()
if not self.require_per_service_call_history_persistence or not history_providers:
return []
conversation_id = (
session.service_session_id
if session and session.service_session_id
else cast(str | None, (options or {}).get("conversation_id") or self.default_options.get("conversation_id"))
)
if service_stores_history:
return []
if conversation_id is not None:
# A live service-managed session id takes precedence over the resolved conversation id.
if session and session.service_session_id:
conversation_id = session.service_session_id
# Without service-side storage the middleware persists locally and drives the function
# loop with a local sentinel, which cannot be reconciled with an existing service-managed
# conversation. When the service stores history, an existing conversation id is expected.
if conversation_id is not None and not service_stores_history:
raise AgentInvalidRequestException(
"require_per_service_call_history_persistence cannot be used "
"with an existing service-managed conversation."
@@ -1167,18 +1177,34 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
input_messages = normalize_messages(messages)
# `store` in runtime or agent options takes precedence over client-level storage
# indicators. An explicit `store=False` forces local (in-memory) history injection,
# even if the client is configured to use service-side storage by default.
store_ = opts.get("store", self.default_options.get("store", getattr(self.client, "STORES_BY_DEFAULT", False)))
# Combine agent-level defaults with runtime options up front so the decisions below read
# `store` from a single place rather than introspecting both dicts. _merge_options applies
# the same precedence used for the actual client call (runtime wins; unset/None falls back
# to the agent default).
effective_options = _merge_options(self.default_options, opts)
# `store` in runtime or agent options takes precedence over the client's default
# storage behavior. An explicit `store=False` forces local (in-memory) history
# injection even when the client stores server-side by default; an explicit
# `store=True` forces service-side storage. A `store=None`/unset value means the
# service falls back to its own default.
explicit_store = effective_options.get("store")
# Internal behavior hint: will the service own history for this run? Only when the
# user left `store` unset do we fall back to the client's STORES_BY_DEFAULT.
service_stores_history = (
explicit_store if explicit_store is not None else getattr(self.client, "STORES_BY_DEFAULT", False)
)
# Resolve conversation_id from the same combined view so an agent-level default is honored
# when the runtime omits it (a live session id still takes precedence below).
effective_conversation_id = effective_options.get("conversation_id")
# Auto-inject InMemoryHistoryProvider when session is provided, no context providers
# registered, and no service-side storage indicators
if (
session is not None
and not self.context_providers
and not session.service_session_id
and not opts.get("conversation_id")
and not store_
and not effective_conversation_id
and not service_stores_history
):
self.context_providers.append(InMemoryHistoryProvider())
@@ -1188,10 +1214,30 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
per_service_call_history_providers = self._resolve_per_service_call_history_providers(
session=active_session,
options=opts,
service_stores_history=bool(store_),
conversation_id=effective_conversation_id,
service_stores_history=service_stores_history,
)
# When require_per_service_call_history_persistence is set together with a
# HistoryProvider, the per-service-call middleware (installed below) always persists
# the provider. ``service_stores_history`` only selects how the middleware behaves:
# - service does not store: the middleware also loads providers and drives the function
# loop with a local sentinel conversation id, or
# - service stores: the middleware skips loading (the service owns history) and simply
# persists each service call while the real conversation id flows through.
# In the service-managed case loading is skipped, so warn for providers that expect to load.
history_providers = self._get_history_providers()
if self.require_per_service_call_history_persistence and history_providers and service_stores_history:
for provider in history_providers:
if provider.load_messages:
logger.warning(
"HistoryProvider '%s' has load_messages=True but the chat client stores history "
"server-side; skipping local history load and relying on the service-managed "
"conversation. Set store=False to load from the provider, or load_messages=False "
"to silence this warning.",
provider.source_id,
)
session_context, chat_options = await self._prepare_session_and_messages(
session=active_session,
input_messages=input_messages,
@@ -1265,8 +1311,8 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
}
if model is not None:
run_opts["model"] = model
# Remove None values and merge with chat_options
run_opts = {k: v for k, v in run_opts.items() if v is not None}
# _merge_options strips unset (None) options, so e.g. an unset `store` is not forwarded
# and the service decides its own default.
co = _merge_options(chat_options, run_opts)
# Build session_messages from session context: context messages + input messages
@@ -1280,6 +1326,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
agent=self,
session=active_session,
providers=per_service_call_history_providers,
service_stores_history=service_stores_history,
)
existing_middleware = effective_client_kwargs.get("middleware")
if isinstance(existing_middleware, Sequence) and not isinstance(existing_middleware, (str, bytes)):
@@ -1319,7 +1366,7 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
"input_messages": input_messages,
"session_messages": session_messages,
"agent_name": agent_name,
"suppress_response_id": bool(per_service_call_history_providers),
"suppress_response_id": bool(per_service_call_history_providers) and not service_stores_history,
"chat_options": co,
"compaction_strategy": compaction_strategy or self.compaction_strategy,
"tokenizer": tokenizer or self.tokenizer,
@@ -1413,11 +1460,15 @@ class RawAgent(BaseAgent, Generic[OptionsCoT]): # type: ignore[misc]
options=options or {},
)
# When per-service-call persistence is enabled, the per-service-call middleware owns
# HistoryProvider loading (it loads locally when the service does not store history, or
# relies on the service when it does), so skip them on the once-per-run before_run path.
per_service_call_history_required = self.require_per_service_call_history_persistence and bool(
self._get_history_providers()
)
# Run before_run providers (forward order, skip HistoryProvider when per-service-call persistence owns history)
# Run before_run providers (forward order, skip HistoryProvider when per-service-call
# persistence owns loading)
for provider in self.context_providers:
if per_service_call_history_required and isinstance(provider, HistoryProvider):
continue
@@ -604,10 +604,13 @@ class BaseChatClient(SerializationMixin, ABC, Generic[OptionsCoT]):
and dict literals are accepted without specialized option typing.
context_providers: Context providers to include during agent invocation.
middleware: List of middleware to intercept agent and function invocations.
require_per_service_call_history_persistence: Whether to require per-service-call
chat history persistence. When enabled, history providers are invoked around
each model call instead of once per ``run()`` when the service is not already
storing history.
require_per_service_call_history_persistence: When enabled (and a HistoryProvider is
present), the provider always persists history after each model call. If the
client does not store history server-side, history providers are also loaded and
injected around each model call; if it does, provider loading is skipped and the
service-managed conversation is the source of truth (persistence still happens
after each model call). When no HistoryProvider is present, this flag has no
effect (no middleware is installed and nothing is persisted).
function_invocation_configuration: Optional function invocation configuration override.
compaction_strategy: Optional agent-level compaction override. When omitted,
client-level compaction defaults remain in effect for each call.
@@ -58,6 +58,7 @@ class ExperimentalFeature(str, Enum):
FOUNDRY_PREVIEW_TOOLS = "FOUNDRY_PREVIEW_TOOLS"
FUNCTIONAL_WORKFLOWS = "FUNCTIONAL_WORKFLOWS"
HARNESS = "HARNESS"
MCP_LONG_RUNNING_TASKS = "MCP_LONG_RUNNING_TASKS"
MCP_SKILLS = "MCP_SKILLS"
PROGRESSIVE_TOOLS = "PROGRESSIVE_TOOLS"
SKILLS = "SKILLS"
File diff suppressed because it is too large Load Diff
@@ -16,6 +16,7 @@ from __future__ import annotations
import asyncio
import copy
import json
import logging
import threading
import uuid
import weakref
@@ -36,6 +37,8 @@ if TYPE_CHECKING:
from ._middleware import MiddlewareTypes
logger = logging.getLogger("agent_framework")
# Registry of known types for state deserialization
_STATE_TYPE_REGISTRY: dict[str, type] = {}
@@ -580,6 +583,7 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
agent: SupportsAgentRun,
session: AgentSession,
providers: Sequence[HistoryProvider],
service_stores_history: bool = False,
) -> None:
"""Initialize the middleware.
@@ -587,10 +591,16 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
agent: The agent that owns the history providers.
session: The active session for the current run.
providers: The history providers participating in per-service-call persistence.
service_stores_history: When True, the chat client stores history server-side. The
middleware then skips loading providers and leaves the real conversation id
untouched, persisting each service call without driving the function loop with a
local sentinel. When False, the middleware loads providers and uses a local
sentinel conversation id so the function loop runs without service-side storage.
"""
self._agent = agent
self._session = session
self._providers = list(providers)
self._service_stores_history = service_stores_history
async def _prepare_service_call_context(self, messages: Sequence[Message]) -> SessionContext:
"""Create a per-call SessionContext and load history providers into it."""
@@ -602,6 +612,9 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
)
for source_id, source_messages in context_messages.items():
service_call_context.extend_messages(source_id, source_messages)
# When the service stores history, it owns loading; the providers are write-only sinks.
if self._service_stores_history:
return service_call_context
for provider in self._providers:
if not provider.load_messages:
continue
@@ -652,17 +665,35 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
response: ChatResponse,
) -> ChatResponse:
"""Persist a model response and apply the local follow-up sentinel when needed."""
if response.conversation_id is not None and not is_local_history_conversation_id(response.conversation_id):
if (
not self._service_stores_history
and response.conversation_id is not None
and not is_local_history_conversation_id(response.conversation_id)
):
raise ChatClientInvalidResponseException(
"require_per_service_call_history_persistence cannot be used "
"when the chat client returns a real conversation_id."
)
# In storing mode the service is expected to echo a conversation id that the next run
# resumes from. If it comes back empty, the provider still captures this turn but there is
# no service id to load from next time, so cross-turn history can be lost silently. Warn
# every time so this uncommon, easy-to-miss failure mode cannot fail quietly.
if self._service_stores_history and response.conversation_id is None:
logger.warning(
"require_per_service_call_history_persistence is enabled with a chat client that "
"stores history server-side, but the client returned no conversation_id; cross-turn "
"history may not resume. Set store=False to load and resume from the HistoryProvider "
"instead."
)
await self._persist_service_call_response(
service_call_context=service_call_context,
response=response,
)
if _response_contains_follow_up_request(response):
# The local sentinel only applies when the service does not store history; when it does,
# the real conversation id already drives function-loop continuation.
if not self._service_stores_history and _response_contains_follow_up_request(response):
response.mark_internal_conversation_id()
response.conversation_id = LOCAL_HISTORY_CONVERSATION_ID
return response
@@ -681,8 +712,12 @@ class PerServiceCallHistoryPersistingMiddleware(ChatMiddleware):
result type for streaming or non-streaming execution.
"""
service_call_context = await self._prepare_service_call_context(context.messages)
context.messages = service_call_context.get_messages(include_input=True)
self._strip_local_conversation_id(context)
# When the service stores history, leave the outgoing messages and the real conversation
# id untouched (pass-through); the middleware only persists. Otherwise reconstruct the
# outgoing messages from the loaded local history and strip the local sentinel.
if not self._service_stores_history:
context.messages = service_call_context.get_messages(include_input=True)
self._strip_local_conversation_id(context)
await call_next()
@@ -2,7 +2,6 @@
from __future__ import annotations
import json
import logging
import sys
import uuid
@@ -12,7 +11,6 @@ from datetime import datetime, timezone
from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload
from .._agents import BaseAgent
from .._serialization import make_json_safe
from .._sessions import (
AgentSession,
ContextProvider,
@@ -30,11 +28,12 @@ from .._types import (
UsageDetails,
add_usage_details,
)
from ..exceptions import AgentInvalidRequestException, AgentInvalidResponseException
from ..exceptions import AgentException, AgentInvalidRequestException, AgentInvalidResponseException
from ._checkpoint import CheckpointStorage
from ._events import (
AGENT_FORWARDED_EVENT_TYPES,
WorkflowEvent,
WorkflowRunState,
)
from ._message_utils import normalize_messages_input
from ._typing_utils import is_instance_of, is_type_compatible
@@ -59,27 +58,24 @@ class WorkflowAgent(BaseAgent):
@dataclass
class RequestInfoFunctionArgs:
request_id: str
data: Any
request_event: WorkflowEvent
def to_dict(self) -> dict[str, Any]:
return {"request_id": self.request_id, "data": make_json_safe(self.data)}
def to_json(self) -> str:
return json.dumps(self.to_dict())
return {"request_id": self.request_id, "request_event": self.request_event.to_dict()}
@classmethod
def from_dict(cls, payload: dict[str, Any]) -> WorkflowAgent.RequestInfoFunctionArgs:
return cls(request_id=payload.get("request_id", ""), data=payload.get("data"))
if "request_id" not in payload or "request_event" not in payload:
raise ValueError(
"Invalid payload for RequestInfoFunctionArgs. 'request_id' and 'request_event' are required."
)
if not payload["request_id"]:
raise ValueError("request_id cannot be empty.")
@classmethod
def from_json(cls, raw: str) -> WorkflowAgent.RequestInfoFunctionArgs:
try:
parsed: Any = json.loads(raw)
except json.JSONDecodeError as exc:
raise ValueError(f"RequestInfoFunctionArgs JSON payload is malformed: {exc}") from exc
if not isinstance(parsed, dict):
raise ValueError("RequestInfoFunctionArgs JSON payload must decode to a mapping")
return cls.from_dict(cast(dict[str, Any], parsed))
return cls(
request_id=payload.get("request_id", ""),
request_event=WorkflowEvent.from_dict(payload.get("request_event", {})),
)
def __init__(
self,
@@ -129,16 +125,11 @@ class WorkflowAgent(BaseAgent):
**kwargs,
)
self._workflow: Workflow = workflow
self._pending_requests: dict[str, WorkflowEvent[Any]] = {}
@property
def workflow(self) -> Workflow:
return self._workflow
@property
def pending_requests(self) -> dict[str, WorkflowEvent[Any]]:
return self._pending_requests
# region Run Methods
@overload
@@ -182,7 +173,7 @@ class WorkflowAgent(BaseAgent):
Args:
messages: The message(s) to send to the workflow. Required for new runs,
should be None when resuming from checkpoint.
could be None if only restoring the underlying workflow from a checkpoint.
Keyword Args:
stream: If True, returns an async iterable of updates. If False (default),
@@ -416,101 +407,79 @@ class WorkflowAgent(BaseAgent):
Yields:
WorkflowEvent objects from the workflow execution.
"""
# Determine the execution mode based on state.
# The streaming flag controls the workflow's internal streaming mode,
# which affects executor behavior (e.g. AgentExecutor emits different event
# types in streaming vs non-streaming mode).
if bool(self.pending_requests):
function_responses = self._process_pending_requests(input_messages)
# Restore the workflow state if a checkpoint is provided
if checkpoint_id is not None:
if checkpoint_storage is None:
raise AgentInvalidRequestException("checkpoint_storage must be provided when checkpoint_id is provided")
logger.debug(f"Restoring workflow from checkpoint {checkpoint_id}")
# Restore the workflow from checkpoint
if streaming:
async for event in self.workflow.run(
responses=function_responses,
stream=True,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
else:
for event in await self.workflow.run(
responses=function_responses,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
elif checkpoint_id is not None:
# Restore the prior workflow state from the checkpoint. Shared
# state (e.g. accumulated conversation history maintained by the
# workflow's executors) survives across turns because Workflow.run
# no longer wipes state per call. Callers who want to deliver a
# new user message after restore should make a second
# `workflow.run(message=...)` call - they are NOT mutually
# exclusive on the same instance, but each must be its own call.
if streaming:
async for event in self.workflow.run(
async for _ in self.workflow.run(
stream=True,
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
):
pass
else:
_ = await self.workflow.run(
checkpoint_id=checkpoint_id,
checkpoint_storage=checkpoint_storage,
)
if not input_messages:
logger.info("No input messages provided; the workflow has been restored to the checkpoint state.")
return
final_state = self._workflow.status
logger.debug(f"Workflow state: {final_state}")
if final_state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
# Extract function responses from input messages, and ensure that
# only function responses are present in messages if there is any
# pending request.
# NOTE: It is possible that some pending requests are not fulfilled,
# and we will let the workflow to handle this -- the agent does not
# have an opinion on this.
function_responses = self._extract_function_responses(input_messages)
if streaming:
async for event in self.workflow.run(
responses=function_responses,
stream=True,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
else:
for event in await self.workflow.run(
checkpoint_id=checkpoint_id,
responses=function_responses,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
elif final_state == WorkflowRunState.IDLE:
if streaming:
async for event in self.workflow.run(
message=input_messages,
stream=True,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
else:
for event in await self.workflow.run(
message=input_messages,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
else:
if streaming:
async for event in self.workflow.run(
message=input_messages,
stream=True,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
else:
for event in await self.workflow.run(
message=input_messages,
checkpoint_storage=checkpoint_storage,
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
):
yield event
raise AgentException(f"The underlying workflow is in an invalid state to restart: {final_state}.")
# endregion Run Methods
def _process_pending_requests(self, input_messages: Sequence[Message]) -> dict[str, Any]:
"""Process pending requests by extracting function responses and updating state.
Args:
input_messages: Input messages that may contain function responses.
Returns:
A dictionary mapping request IDs to their response data.
"""
logger.info(f"Continuing workflow to address {len(self.pending_requests)} requests")
# Extract function responses from input messages, and ensure that
# only function responses are present in messages if there is any
# pending request.
function_responses = self._extract_function_responses(input_messages)
# Pop pending requests if fulfilled.
for request_id in list(self.pending_requests.keys()):
if request_id in function_responses:
self.pending_requests.pop(request_id)
# NOTE: It is possible that some pending requests are not fulfilled,
# and we will let the workflow to handle this -- the agent does not
# have an opinion on this.
return function_responses
def _convert_workflow_events_to_agent_response(
self,
response_id: str,
@@ -528,10 +497,10 @@ class WorkflowAgent(BaseAgent):
for output_event in output_events:
if output_event.type == "request_info":
function_call, approval_request = self._process_request_info_event(output_event)
request_content = self._process_request_info_event(output_event)
messages.append(
Message(
contents=[function_call, approval_request],
contents=[request_content],
role="assistant",
author_name=output_event.source_executor_id,
message_id=str(uuid.uuid4()),
@@ -598,38 +567,6 @@ class WorkflowAgent(BaseAgent):
raw_representation=raw_representations,
)
def _process_request_info_event(
self,
event: WorkflowEvent[Any],
) -> tuple[Content, Content]:
"""Convert a request_info event to FunctionCallContent and FunctionApprovalRequestContent.
Args:
event: A WorkflowEvent with type='request_info'.
Returns:
A tuple of (FunctionCallContent, FunctionApprovalRequestContent).
"""
request_id = event.request_id
if not request_id:
raise ValueError("request_info event must have a request_id")
self.pending_requests[request_id] = event
args = self.RequestInfoFunctionArgs(request_id=request_id, data=event.data).to_dict()
function_call = Content.from_function_call(
call_id=request_id,
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=args,
)
approval_request = Content.from_function_approval_request(
id=request_id,
function_call=function_call,
additional_properties={"request_id": request_id},
)
return function_call, approval_request
def _convert_workflow_event_to_agent_response_updates(
self,
response_id: str,
@@ -731,85 +668,72 @@ class WorkflowAgent(BaseAgent):
]
if event.type == "request_info":
# Store the pending request for later correlation
request_id = event.request_id
if not request_id:
raise ValueError("request_info event must have a request_id")
self.pending_requests[request_id] = event
args = self.RequestInfoFunctionArgs(request_id=request_id, data=event.data).to_dict()
function_call = Content.from_function_call(
call_id=request_id,
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=args,
)
approval_request = Content.from_function_approval_request(
id=request_id,
function_call=function_call,
additional_properties={"request_id": request_id},
)
request_content = self._process_request_info_event(event)
return [
AgentResponseUpdate(
contents=[function_call, approval_request],
contents=[request_content],
role="assistant",
author_name=self.name,
response_id=response_id,
message_id=str(uuid.uuid4()),
created_at=datetime.now(tz=timezone.utc).strftime("%Y-%m-%dT%H:%M:%S.%fZ"),
raw_representation=event,
)
]
# Ignore workflow-internal events
return []
def _process_request_info_event(
self,
event: WorkflowEvent[Any],
) -> Content:
"""Convert a request_info event to FunctionApprovalRequestContent.
Args:
event: A WorkflowEvent with type='request_info'.
Returns:
A content object representing the request info. The content can be a `function_approval_request`
or a `function_call` depending on the structure of the event data.
Note:
If the event data is already a FunctionApprovalRequestContent, it will be returned as-is.
"""
if isinstance(event.data, Content) and event.data.user_input_request:
# Return the event data as-is if it's already a properly formed FunctionApprovalRequestContent
return event.data
request_id = event.request_id
args = self.RequestInfoFunctionArgs(request_id=request_id, request_event=event).to_dict()
return Content.from_function_call(
call_id=request_id,
name=self.REQUEST_INFO_FUNCTION_NAME,
arguments=args,
)
def _extract_function_responses(self, input_messages: Sequence[Message]) -> dict[str, Any]:
"""Extract function responses from input messages."""
"""Extract function responses from input messages.
The responses are for pending requests that the workflow is waiting on, and
will be passed to the workflow. The pending requests are processed to either
`function_approval_request` or `function_call` content by `_process_request_info_event`.
"""
function_responses: dict[str, Any] = {}
for message in input_messages:
for content in message.contents:
if content.type == "function_approval_response":
# Parse the function arguments to recover request payload
arguments_payload = content.function_call.arguments # type: ignore[attr-defined, union-attr]
if isinstance(arguments_payload, str):
try:
parsed_args = self.RequestInfoFunctionArgs.from_json(arguments_payload)
except ValueError as exc:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent arguments must decode to a mapping."
) from exc
elif isinstance(arguments_payload, dict):
parsed_args = self.RequestInfoFunctionArgs.from_dict(arguments_payload)
else:
raise AgentInvalidResponseException(
"FunctionApprovalResponseContent arguments must be a mapping or JSON string."
)
request_id = parsed_args.request_id or content.id # type: ignore[attr-defined]
if not content.approved: # type: ignore[attr-defined]
raise AgentInvalidResponseException(f"Request '{request_id}' was not approved by the caller.")
if request_id in self.pending_requests:
function_responses[request_id] = parsed_args.data
elif bool(self.pending_requests):
raise AgentInvalidRequestException(
"Only responses for pending requests are allowed when there are outstanding approvals."
)
request_id: str = content.id # type: ignore[assignment]
function_responses[request_id] = content
elif content.type == "function_result":
request_id = content.call_id # type: ignore[attr-defined]
if request_id in self.pending_requests:
response_data = content.result if hasattr(content, "result") else str(content) # type: ignore[attr-defined]
function_responses[request_id] = response_data
elif bool(self.pending_requests):
raise AgentInvalidRequestException(
"Only function responses for pending requests are allowed while requests are outstanding."
)
response_data = content.result if hasattr(content, "result") else str(content) # type: ignore[attr-defined]
function_responses[content.call_id] = response_data # type: ignore
else:
if bool(self.pending_requests):
raise AgentInvalidResponseException(
"Unexpected content type while awaiting request info responses."
)
raise AgentInvalidResponseException(
"Unexpected content type while awaiting request info responses."
)
return function_responses
def _extract_contents(self, data: Any) -> list[Content]:
@@ -429,15 +429,30 @@ class AgentExecutor(Executor):
function_invocation_kwargs=function_invocation_kwargs,
client_kwargs=client_kwargs,
)
await ctx.yield_output(response)
# Handle any user input requests
if response.user_input_requests:
user_input_request_count = len(response.user_input_requests)
total_message_content_count = sum(len(msg.contents) for msg in response.messages)
if user_input_request_count != total_message_content_count:
logger.warning(
"Response %s contains %d user input requests but total message contents are %d. "
"This indicates the response contains both user input requests and message contents. "
"Double check if this is the intended behavior, as non user input request contents in "
"this response will not be emitted.",
response.response_id,
user_input_request_count,
total_message_content_count,
)
for user_input_request in response.user_input_requests:
self._pending_agent_requests[user_input_request.id] = user_input_request # type: ignore[index]
await ctx.request_info(user_input_request, Content)
await ctx.request_info(user_input_request, Content, request_id=user_input_request.id)
return None
# Only yield output if the response is complete and not waiting for user input.
# This is to avoid emitting two events of different types ('output' and 'request_info')
# that carry the same payload.
await ctx.yield_output(response)
return response
async def _run_agent_streaming(self, ctx: WorkflowContext[Never, AgentResponseUpdate]) -> AgentResponse | None:
@@ -472,9 +487,25 @@ class AgentExecutor(Executor):
)
async for update in stream:
updates.append(update)
await ctx.yield_output(update)
if update.user_input_requests:
user_input_request_count = len(update.user_input_requests)
total_message_content_count = len(update.contents)
if user_input_request_count != total_message_content_count:
logger.warning(
"Response update %s contains %d user input requests but total message contents are %d. "
"This indicates the response update contains both user input requests and message contents. "
"Double check if this is the intended behavior, as non user input request contents will "
"not be emitted.",
update.response_id,
user_input_request_count,
total_message_content_count,
)
streamed_user_input_requests.extend(update.user_input_requests)
else:
# Only yield output events for updates that do not contain user input requests.
# This is to avoid emitting two events of different types ('output' and 'request_info')
# that carry the same payload.
await ctx.yield_output(update)
# Prefer stream finalization when available so result hooks run
# (e.g., thread conversation updates). Fall back to reconstructing from updates
@@ -509,7 +540,7 @@ class AgentExecutor(Executor):
if user_input_requests:
for user_input_request in user_input_requests:
self._pending_agent_requests[user_input_request.id] = user_input_request # type: ignore[index]
await ctx.request_info(user_input_request, Content)
await ctx.request_info(user_input_request, Content, request_id=user_input_request.id)
return None
return response
@@ -360,6 +360,22 @@ class Workflow(DictConvertible):
# Flag to prevent concurrent workflow executions
self._is_running = False
# Current run-level status of this workflow instance. Updated in lockstep with
# the status events emitted from `_run_workflow_with_tracing`. Defaults to IDLE
# for a freshly built workflow that has not yet been run.
self._status: WorkflowRunState = WorkflowRunState.IDLE
@property
def status(self) -> WorkflowRunState:
"""Return the current run-level status of this workflow instance.
Mirrors the most recent status event emitted by the workflow. Safe to read at
any time: workflows run on a single asyncio event loop, and the underlying
attribute is a single enum reference whose assignment is atomic under the
CPython GIL, so no locking is required.
"""
return self._status
def _ensure_not_running(self) -> None:
"""Ensure the workflow is not already running."""
if self._is_running:
@@ -513,8 +529,9 @@ class Workflow(DictConvertible):
with _framework_event_origin():
started = WorkflowEvent.started()
yield started # noqa: RUF070
self._status = WorkflowRunState.IN_PROGRESS
with _framework_event_origin():
in_progress = WorkflowEvent.status(WorkflowRunState.IN_PROGRESS)
in_progress = WorkflowEvent.status(self._status)
yield in_progress # noqa: RUF070
# Per-run reset for fresh-message runs only. We deliberately
@@ -569,17 +586,20 @@ class Workflow(DictConvertible):
if event.type == "request_info" and not emitted_in_progress_pending:
emitted_in_progress_pending = True
self._status = WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
with _framework_event_origin():
pending_status = WorkflowEvent.status(WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS)
pending_status = WorkflowEvent.status(self._status)
yield pending_status # noqa: RUF070
# Workflow runs until idle - emit final status based on whether requests are pending
if saw_request:
self._status = WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
with _framework_event_origin():
terminal_status = WorkflowEvent.status(WorkflowRunState.IDLE_WITH_PENDING_REQUESTS)
terminal_status = WorkflowEvent.status(self._status)
yield terminal_status
else:
self._status = WorkflowRunState.IDLE
with _framework_event_origin():
terminal_status = WorkflowEvent.status(WorkflowRunState.IDLE)
terminal_status = WorkflowEvent.status(self._status)
yield terminal_status
span.add_event(OtelAttr.WORKFLOW_COMPLETED)
@@ -593,6 +613,7 @@ class Workflow(DictConvertible):
with _framework_event_origin():
failed_event = WorkflowEvent.failed(details)
yield failed_event # noqa: RUF070
self._status = WorkflowRunState.FAILED
with _framework_event_origin():
failed_status = WorkflowEvent.status(WorkflowRunState.FAILED)
yield failed_status # noqa: RUF070
@@ -80,6 +80,7 @@ __all__ = [
"EmbeddingTelemetryLayer",
"OtelAttr",
"configure_otel_providers",
"create_mcp_client_span",
"create_metric_views",
"create_resource",
"disable_instrumentation",
@@ -87,6 +88,7 @@ __all__ = [
"enable_sensitive_telemetry",
"get_meter",
"get_tracer",
"set_mcp_span_error",
]
@@ -110,7 +112,6 @@ INNER_ACCUMULATED_USAGE: Final[contextvars.ContextVar[UsageDetails | None]] = co
"inner_accumulated_usage", default=None
)
OTEL_METRICS: Final[str] = "__otel_metrics__"
TOKEN_USAGE_BUCKET_BOUNDARIES: Final[tuple[float, ...]] = (
1,
@@ -292,6 +293,14 @@ class OtelAttr(str, Enum):
AGENT_CREATE_OPERATION = "create_agent"
AGENT_INVOKE_OPERATION = "invoke_agent"
# MCP attributes (https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/)
MCP_METHOD_NAME = "mcp.method.name"
MCP_PROTOCOL_VERSION = "mcp.protocol.version"
MCP_SESSION_ID = "mcp.session.id"
PROMPT_NAME = "gen_ai.prompt.name"
NETWORK_TRANSPORT = "network.transport"
NETWORK_PROTOCOL_NAME = "network.protocol.name"
# Agent Framework specific attributes
MEASUREMENT_FUNCTION_TAG_NAME = "agent_framework.function.name"
MEASUREMENT_FUNCTION_INVOCATION_DURATION = "agent_framework.function.invocation.duration"
@@ -2013,6 +2022,61 @@ def get_function_span(
)
# region MCP span helpers
@contextlib.contextmanager
def create_mcp_client_span(
method_name: str,
target: str | None = None,
attributes: dict[str, Any] | None = None,
) -> Generator[trace.Span, Any, Any]:
"""Create an MCP client span per OTel MCP semantic conventions.
Span name follows the format ``{mcp.method.name} {target}`` when a target
is available, otherwise just ``{mcp.method.name}``.
See: https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/#client
Args:
method_name: The MCP method name (e.g. ``initialize``, ``tools/call``).
target: Optional low-cardinality target (tool name, prompt name).
attributes: Additional span attributes.
"""
span_name = f"{method_name} {target}" if target else method_name
attrs: dict[str, Any] = {OtelAttr.MCP_METHOD_NAME: method_name}
if attributes:
attrs.update(attributes)
tracer = get_tracer() if OBSERVABILITY_SETTINGS.ENABLED else trace.NoOpTracer()
span = tracer.start_span(span_name, kind=trace.SpanKind.CLIENT, attributes=attrs)
with trace.use_span(
span=span,
end_on_exit=True,
record_exception=True,
set_status_on_exception=True,
) as current_span:
yield current_span
def set_mcp_span_error(
span: trace.Span,
error_type: str,
description: str | None = None,
) -> None:
"""Set error status and ``error.type`` on an MCP span.
Args:
span: The span to mark as errored.
error_type: The error type string (e.g. ``tool_error``, exception class name).
description: Optional description (e.g. JSON-RPC error message).
"""
span.set_attribute(OtelAttr.ERROR_TYPE, error_type)
span.set_status(trace.StatusCode.ERROR, description=description)
# endregion
@contextlib.contextmanager
def _activate_span(span: trace.Span) -> Generator[None]:
"""Attach ``span`` as the current span in the OpenTelemetry context.
+472 -2
View File
@@ -3,6 +3,7 @@
import contextlib
import inspect
import json
import logging
from collections.abc import AsyncIterable, Awaitable, Callable, MutableSequence, Sequence
from typing import Any, cast
from unittest.mock import AsyncMock, MagicMock, patch
@@ -42,6 +43,8 @@ from agent_framework._mcp import MCPTool, _build_prefixed_mcp_name, _normalize_m
from agent_framework._middleware import FunctionInvocationContext
from agent_framework.exceptions import AgentInvalidRequestException, ChatClientInvalidResponseException
from .conftest import MockBaseChatClient
class _FixedTokenizer:
def __init__(self, token_count: int) -> None:
@@ -609,6 +612,7 @@ async def test_streaming_per_service_call_persistence_hides_response_id_from_aft
async def test_per_service_call_persistence_uses_real_service_storage_when_client_stores_by_default(
chat_client_base: SupportsChatGetResponse,
caplog: pytest.LogCaptureFixture,
) -> None:
provider = _RecordingHistoryProvider()
@@ -649,15 +653,22 @@ async def test_per_service_call_persistence_uses_real_service_storage_when_clien
require_per_service_call_history_persistence=True,
)
result = await agent.run("What's the weather in Seattle?", session=session)
with caplog.at_level(logging.WARNING, logger="agent_framework"):
result = await agent.run("What's the weather in Seattle?", session=session)
provider_state = session.state[provider.source_id]
assert result.text == "It is sunny in Seattle."
assert result.response_id == "resp_call_2"
assert chat_client_base.call_count == 2
# The service owns the conversation, so the provider never loads (issue #5798).
assert "get_call_count" not in provider_state
assert "save_call_count" not in provider_state
# Persistence is owned by the per-service-call middleware: it persists once per service call
# (issue #5798: the provider must never be silently bypassed when the service stores history).
# This run makes two service calls (function call + final answer), so it persists twice.
assert provider_state["save_call_count"] == 2
# load_messages=True while the service stores history surfaces a warning.
assert any("load_messages" in record.message for record in caplog.records)
assert session.service_session_id == "resp_service_managed"
@@ -1996,6 +2007,19 @@ def test_merge_options_none_values_ignored():
assert result["key2"] == "value2"
def test_merge_options_drops_none_base_values():
"""Test _merge_options strips None values so unset options are never forwarded."""
base = {"store": None, "temperature": 0.5}
override = {"top_p": 0.9}
result = _merge_options(base, override)
# An unset base value (e.g. store=None from default_options) must not survive the merge.
assert "store" not in result
assert result["temperature"] == 0.5
assert result["top_p"] == 0.9
def test_merge_options_runtime_model_overrides_default_model() -> None:
"""Test _merge_options lets a runtime model override a default model."""
result = _merge_options({"model": "default-model"}, {"model": "runtime-model"})
@@ -2658,3 +2682,449 @@ async def test_as_tool_raises_on_user_input_request(client: SupportsChatGetRespo
assert len(exc_info.value.contents) == 1
assert exc_info.value.contents[0].type == "oauth_consent_request"
assert exc_info.value.contents[0].consent_link == "https://login.microsoftonline.com/consent"
# region Per-service-call history persistence scenario matrix
#
# The driving field is ``require_per_service_call_history_persistence``. Every scenario runs a
# single agent run that makes **two service calls** -- a function call followed by a final
# completion -- so the *timing* of persistence is observable:
#
# * When the flag is ``True``, the per-service-call middleware persists the provider **after each
# service call**. So the function-call turn is already saved by the time the second (final)
# service call starts. This holds regardless of whether the chat client stores history
# server-side (the bug in issue #5798 was that a storing client silently bypassed persistence).
# * When the flag is ``False``, the provider persists **once, at the end of the run** -- nothing is
# saved between the two service calls.
#
# ``SpyChatClient.saves_before_call`` records ``provider.save_calls`` at the start of every service
# call, so ``[0, 1]`` means "the function-call turn was persisted before the final call" and
# ``[0, 0]`` means "no persistence happened mid-run". The client's ``store`` / ``STORES_BY_DEFAULT``
# only selects *how* the middleware behaves -- never *whether* the provider persists.
_PSC_SERVICE_CONVERSATION_ID = "svc-conversation"
_psc_stream_params = pytest.mark.parametrize("stream", [False, True], ids=["sync", "stream"])
@tool(name="lookup_weather", approval_mode="never_require")
def _psc_lookup_weather(location: str) -> str:
return f"Weather in {location}: sunny"
def _psc_function_call_script() -> list[tuple[str, ...]]:
"""A fresh function-call-then-final-completion script (the client mutates it)."""
return [
("call", "call_1", "lookup_weather", '{"location": "Seattle"}'),
("text", "It is sunny in Seattle."),
]
class _PscSpyHistoryProvider(HistoryProvider):
"""In-memory history provider that records load/save calls for assertions."""
def __init__(self, source_id: str = "spy_history", **kwargs: Any) -> None:
super().__init__(source_id, **kwargs)
self._messages: list[Message] = []
self.get_calls: int = 0
self.save_calls: int = 0
self.saved_batches: list[list[Message]] = []
async def get_messages(
self, session_id: str | None, *, state: dict[str, Any] | None = None, **kwargs: Any
) -> list[Message]:
self.get_calls += 1
return list(self._messages)
async def save_messages(
self,
session_id: str | None,
messages: Sequence[Message],
*,
state: dict[str, Any] | None = None,
**kwargs: Any,
) -> None:
self.save_calls += 1
self.saved_batches.append(list(messages))
self._messages.extend(messages)
@property
def stored_messages(self) -> list[Message]:
return list(self._messages)
class _PscSpyChatClient(MockBaseChatClient):
"""Chat client that scripts a function-call/final-completion sequence.
It records, at the start of each service call, how many provider saves have already happened
(``saves_before_call``), what messages it received, and what options it saw. When the effective
``store`` is truthy it returns a stable ``conversation_id`` to mimic a server-managed
conversation, so the framework propagates ``session.service_session_id``.
"""
def __init__(
self,
*,
provider: _PscSpyHistoryProvider,
stores_by_default: bool = False,
script: list[tuple[str, ...]] | None = None,
echo_conversation_id: bool = True,
**kwargs: Any,
) -> None:
super().__init__(**kwargs)
self.STORES_BY_DEFAULT = stores_by_default # type: ignore[attr-defined]
self._provider = provider
self._script = list(script) if script is not None else [("text", "ok")]
self._echo_conversation_id = echo_conversation_id
self.received_messages: list[list[Message]] = []
self.received_options: list[dict[str, Any]] = []
self.saves_before_call: list[int] = []
def _effective_store(self, options: dict[str, Any]) -> bool:
store = options.get("store")
if store is None:
return bool(self.STORES_BY_DEFAULT)
return bool(store)
def _next_contents(self) -> list[Content]:
turn = self._script.pop(0) if self._script else ("text", "ok")
if turn[0] == "call":
_, call_id, name, args = turn
return [Content.from_function_call(call_id=call_id, name=name, arguments=args)]
return [Content.from_text(turn[1])]
def _inner_get_response( # type: ignore[override]
self,
*,
messages: MutableSequence[Message],
stream: bool,
options: dict[str, Any],
**kwargs: Any,
) -> Awaitable[ChatResponse] | ResponseStream[ChatResponseUpdate, ChatResponse]:
self.received_messages.append(list(messages))
self.received_options.append(dict(options))
self.saves_before_call.append(self._provider.save_calls)
store_and_echo = self._effective_store(options) and self._echo_conversation_id
conv_id = _PSC_SERVICE_CONVERSATION_ID if store_and_echo else None
contents = self._next_contents()
if stream:
async def _stream() -> AsyncIterable[ChatResponseUpdate]:
self.call_count += 1
yield ChatResponseUpdate(
contents=contents,
role="assistant",
finish_reason="stop",
conversation_id=conv_id,
)
def _finalize(updates: Sequence[ChatResponseUpdate]) -> ChatResponse:
response = ChatResponse.from_updates(updates, output_format_type=options.get("response_format"))
if conv_id:
response.conversation_id = conv_id
return response
return ResponseStream(_stream(), finalizer=_finalize)
async def _get() -> ChatResponse:
self.call_count += 1
return ChatResponse(
messages=Message(role="assistant", contents=contents),
conversation_id=conv_id,
)
return _get()
def _psc_build_agent(
client: _PscSpyChatClient,
provider: _PscSpyHistoryProvider,
*,
require_per_service_call_history_persistence: bool,
default_options: dict[str, Any] | None = None,
) -> Agent:
kwargs: dict[str, Any] = {}
if default_options is not None:
kwargs["default_options"] = default_options
return Agent(
client=client,
tools=[_psc_lookup_weather],
context_providers=[provider],
require_per_service_call_history_persistence=require_per_service_call_history_persistence,
**kwargs,
)
async def _psc_run(agent: Agent, text: str, session: AgentSession, *, stream: bool) -> str:
if stream:
chunks: list[str] = []
async for update in agent.run(text, session=session, stream=True):
chunks.append(update.text or "")
return "".join(chunks)
result = await agent.run(text, session=session)
return result.text
# driver=True (the contract under test): persistence happens per service call
@_psc_stream_params
async def test_psc_flag_on_store_false_persists_after_each_service_call(stream: bool) -> None:
"""Mode A (flag on, service does not store): function-call turn is persisted before the final call."""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=False, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
text = await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert text == "It is sunny in Seattle."
# Two service calls: function call, then final completion.
assert client.call_count == 2
# The contract: the function-call turn was persisted *before* the second service call started.
assert client.saves_before_call == [0, 1]
assert provider.save_calls == 2
# Mode A loads local history (the middleware injects it before each service call).
assert provider.get_calls >= 1
# No service-side storage, so no conversation id is propagated.
assert session.service_session_id is None
@_psc_stream_params
async def test_psc_flag_on_stores_by_default_persists_after_each_service_call(
stream: bool, caplog: pytest.LogCaptureFixture
) -> None:
"""Mode B (flag on, service stores by default): still persists per service call, but skips load (issue #5798)."""
provider = _PscSpyHistoryProvider() # load_messages=True by default
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
with caplog.at_level(logging.WARNING, logger="agent_framework"):
text = await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert text == "It is sunny in Seattle."
assert client.call_count == 2
# The invariant the bug violated: persistence still happens per service call when the service stores.
assert client.saves_before_call == [0, 1]
assert provider.save_calls == 2
# The service owns loading, so the provider is never asked to load.
assert provider.get_calls == 0
# A warning surfaces the bypassed load (load_messages=True).
assert any("load_messages" in record.message for record in caplog.records)
# The real service conversation id propagates to the session.
assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
@_psc_stream_params
async def test_psc_flag_on_store_only_provider_no_load_no_warning(
stream: bool, caplog: pytest.LogCaptureFixture
) -> None:
"""Mode B with a store-only provider (load_messages=False): persists per call, no load, no warning."""
provider = _PscSpyHistoryProvider(load_messages=False)
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
with caplog.at_level(logging.WARNING, logger="agent_framework"):
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert client.saves_before_call == [0, 1]
assert provider.save_calls == 2
assert provider.get_calls == 0
assert not any("load_messages" in record.message for record in caplog.records)
@_psc_stream_params
async def test_psc_flag_on_store_false_override_behaves_as_mode_a(stream: bool) -> None:
"""Flag on + storing client but store=False override: falls back to Mode A (local, per call)."""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(
client, provider, require_per_service_call_history_persistence=True, default_options={"store": False}
)
session = agent.create_session()
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert client.saves_before_call == [0, 1]
assert provider.save_calls == 2
assert provider.get_calls >= 1
# store=False forces local handling, so no real service conversation id.
assert session.service_session_id is None
@_psc_stream_params
async def test_psc_flag_on_store_none_treated_as_absent(stream: bool, caplog: pytest.LogCaptureFixture) -> None:
"""Flag on + storing client + explicit store=None: None is "unset", so the storing default applies (Mode B)."""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(
client, provider, require_per_service_call_history_persistence=True, default_options={"store": None}
)
session = agent.create_session()
with caplog.at_level(logging.WARNING, logger="agent_framework"):
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert client.saves_before_call == [0, 1]
assert provider.save_calls == 2
assert provider.get_calls == 0
assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
assert any("load_messages" in record.message for record in caplog.records)
# store=None must not be forwarded to the client; the service decides its own default.
assert all("store" not in options for options in client.received_options)
@_psc_stream_params
async def test_psc_flag_on_respects_store_outputs_flag(stream: bool) -> None:
"""Flag on: the provider's store_inputs/store_outputs flags still apply per service call."""
provider = _PscSpyHistoryProvider(store_inputs=True, store_outputs=False)
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert provider.save_calls == 2
# Outputs disabled, so no assistant/tool-call messages were stored, only user/tool inputs.
assert provider.stored_messages
assert all(message.role != "assistant" for message in provider.stored_messages)
# driver=False (control): persistence happens once, at the end of the run
@_psc_stream_params
async def test_psc_flag_off_store_false_persists_once_at_end(stream: bool) -> None:
"""Flag off + non-storing client: nothing is persisted mid-run; one save at the end."""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=False, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=False)
session = agent.create_session()
text = await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert text == "It is sunny in Seattle."
assert client.call_count == 2
# The control contract: no save happened between the function call and the final completion.
assert client.saves_before_call == [0, 0]
assert provider.save_calls == 1
@_psc_stream_params
async def test_psc_flag_off_stores_by_default_persists_once_at_end(stream: bool) -> None:
"""Flag off + storing client: once-per-run persistence, and the service conversation id propagates."""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=False)
session = agent.create_session()
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert client.saves_before_call == [0, 0]
assert provider.save_calls == 1
assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
@_psc_stream_params
async def test_psc_flag_on_storing_with_existing_conversation_id_does_not_raise(stream: bool) -> None:
"""Allow side of the guard: flag on + storing client + an existing conversation_id resumes (no raise).
The non-storing path raises on an existing service-managed conversation id, but with a storing
client the run must proceed and the service conversation id must propagate to the session.
"""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
if stream:
chunks: list[str] = []
async for update in agent.run(
"What's the weather in Seattle?",
session=session,
stream=True,
options={"conversation_id": "existing_conversation"},
):
chunks.append(update.text or "")
text = "".join(chunks)
else:
result = await agent.run(
"What's the weather in Seattle?",
session=session,
options={"conversation_id": "existing_conversation"},
)
text = result.text
assert text == "It is sunny in Seattle."
# Persistence still happens per service call, and the real service id propagates to the session.
assert provider.save_calls == 2
assert provider.get_calls == 0
assert session.service_session_id == _PSC_SERVICE_CONVERSATION_ID
@_psc_stream_params
async def test_psc_flag_on_storing_two_runs_same_session(stream: bool) -> None:
"""Storing mode across two runs on one session: persistence keeps happening, id is stable, no load.
The second run exercises the precedence branch where the session already carries a
service_session_id, which must continue to skip provider loading and keep persisting.
"""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(provider=provider, stores_by_default=True, script=_psc_function_call_script())
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
assert provider.save_calls == 2
assert provider.get_calls == 0
first_run_service_id = session.service_session_id
assert first_run_service_id == _PSC_SERVICE_CONVERSATION_ID
# Reset the scripted client for a second run on the same session.
client._script = _psc_function_call_script()
client.call_count = 0
client.saves_before_call = []
await _psc_run(agent, "And in Portland?", session, stream=stream)
# Persistence keeps happening on the second run (two more saves), still per service call.
assert client.saves_before_call == [2, 3]
assert provider.save_calls == 4
# Loading stays skipped and the service conversation id stays stable across runs.
assert provider.get_calls == 0
assert session.service_session_id == first_run_service_id
@_psc_stream_params
async def test_psc_flag_on_storing_without_conversation_id_warns_every_call(
stream: bool, caplog: pytest.LogCaptureFixture
) -> None:
"""Storing mode but the client returns no conversation_id: warn on every service call.
Without an echoed conversation id the next run has nothing to resume from, so cross-turn
history can be lost silently. The warning fires per service call (no dedup) so the uncommon
failure mode cannot pass unnoticed.
"""
provider = _PscSpyHistoryProvider()
client = _PscSpyChatClient(
provider=provider,
stores_by_default=True,
script=_psc_function_call_script(),
echo_conversation_id=False,
)
agent = _psc_build_agent(client, provider, require_per_service_call_history_persistence=True)
session = agent.create_session()
with caplog.at_level(logging.WARNING, logger="agent_framework"):
await _psc_run(agent, "What's the weather in Seattle?", session, stream=stream)
# Persistence still happens, but no service id is captured to resume from.
assert provider.save_calls == 2
assert session.service_session_id is None
# Two service calls -> the warning is emitted twice (one per call, not deduped).
missing_id_warnings = [r for r in caplog.records if "returned no conversation_id" in r.message]
assert len(missing_id_warnings) == 2
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,376 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for MCP client span instrumentation per OTel GenAI Semantic Conventions.
See: https://opentelemetry.io/docs/specs/semconv/gen-ai/mcp/#client
"""
from __future__ import annotations
from typing import Any
from unittest.mock import AsyncMock, Mock
import pytest
from mcp import types
from mcp.shared.exceptions import McpError
from mcp.types import ErrorData
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
from opentelemetry.trace import SpanKind, StatusCode
from agent_framework import MCPStdioTool, MCPStreamableHTTPTool, MCPWebsocketTool
from agent_framework._mcp import MCPTool
from agent_framework.exceptions import ToolExecutionException
from agent_framework.observability import OtelAttr
# region helpers
def _make_connected_mcp_tool(
name: str = "test-mcp",
*,
supports_tools: bool = True,
supports_prompts: bool = True,
) -> MCPTool:
"""Create an MCPTool with a mocked session, ready for testing."""
tool = MCPTool(name=name)
tool.session = AsyncMock()
tool.is_connected = True
tool._supports_tools = supports_tools
tool._supports_prompts = supports_prompts
tool.load_tools_flag = True
tool.load_prompts_flag = True
return tool
def _make_tool_list_result(
tools: list[dict[str, Any]] | None = None,
) -> Mock:
"""Create a mock ListToolsResult."""
if tools is None:
tools = [{"name": "get-weather", "description": "Get weather", "inputSchema": {"type": "object"}}]
result = Mock()
result.tools = [
types.Tool(name=t["name"], description=t.get("description", ""), inputSchema=t.get("inputSchema", {}))
for t in tools
]
result.nextCursor = None
return result
def _make_prompt_list_result(
prompts: list[dict[str, Any]] | None = None,
) -> Mock:
"""Create a mock ListPromptsResult."""
if prompts is None:
prompts = [{"name": "analyze-code", "description": "Analyze code"}]
result = Mock()
result.prompts = [
types.Prompt(name=p["name"], description=p.get("description", ""), arguments=None) for p in prompts
]
result.nextCursor = None
return result
def _make_call_tool_result(text: str = "result", is_error: bool = False) -> Mock:
"""Create a mock CallToolResult."""
result = Mock()
result.isError = is_error
result.content = [types.TextContent(type="text", text=text)]
return result
def _make_get_prompt_result(text: str = "prompt result") -> types.GetPromptResult:
"""Create a mock GetPromptResult."""
return types.GetPromptResult(
description="test prompt",
messages=[
types.PromptMessage(
role="user",
content=types.TextContent(type="text", text=text),
)
],
)
# endregion
# region initialize span
async def test_mcp_initialize_span(span_exporter: InMemorySpanExporter):
"""session.initialize() should produce an MCP CLIENT span named 'initialize'."""
tool = MCPTool(name="test-server")
mock_session_cls = AsyncMock()
init_result = Mock()
init_result.capabilities = None
init_result.protocolVersion = "2025-06-18"
mock_session_cls.initialize = AsyncMock(return_value=init_result)
# Create a mock transport context manager
mock_transport = AsyncMock()
mock_transport.__aenter__ = AsyncMock(return_value=(Mock(), Mock()))
mock_transport.__aexit__ = AsyncMock(return_value=False)
# Mock get_mcp_client and the session creation
tool.session = None
tool.load_tools_flag = False
tool.load_prompts_flag = False
span_exporter.clear()
with pytest.MonkeyPatch.context() as m:
m.setattr(tool, "get_mcp_client", lambda: mock_transport)
async def patched_connect(self_: Any, *, reset: bool = False, load_configured: bool = True) -> None:
# Simulate _connect_on_owner: create initialize span and call session.initialize()
from agent_framework._mcp import create_mcp_client_span
from agent_framework.observability import OtelAttr
with create_mcp_client_span("initialize", attributes=self_._mcp_base_span_attributes()) as init_span:
result = await mock_session_cls.initialize()
protocol_version = getattr(result, "protocolVersion", None)
if protocol_version:
init_span.set_attribute(OtelAttr.MCP_PROTOCOL_VERSION, protocol_version)
self_.session = mock_session_cls
self_.is_connected = True
m.setattr(MCPTool, "_connect_on_owner", patched_connect)
await tool.connect()
mock_session_cls.initialize.assert_awaited_once()
spans = span_exporter.get_finished_spans()
init_spans = [s for s in spans if s.name == "initialize"]
assert len(init_spans) == 1
span = init_spans[0]
assert span.kind == SpanKind.CLIENT
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "initialize"
assert span.attributes.get(OtelAttr.MCP_PROTOCOL_VERSION) == "2025-06-18"
# endregion
# region tools/list span
async def test_mcp_tools_list_span(span_exporter: InMemorySpanExporter):
"""session.list_tools() should produce an MCP CLIENT span named 'tools/list'."""
tool = _make_connected_mcp_tool()
tool.session.list_tools = AsyncMock(return_value=_make_tool_list_result())
span_exporter.clear()
await tool.load_tools()
spans = span_exporter.get_finished_spans()
list_spans = [s for s in spans if s.name == "tools/list"]
assert len(list_spans) == 1
span = list_spans[0]
assert span.kind == SpanKind.CLIENT
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "tools/list"
# endregion
# region prompts/list span
async def test_mcp_prompts_list_span(span_exporter: InMemorySpanExporter):
"""session.list_prompts() should produce an MCP CLIENT span named 'prompts/list'."""
tool = _make_connected_mcp_tool()
tool.session.list_prompts = AsyncMock(return_value=_make_prompt_list_result())
span_exporter.clear()
await tool.load_prompts()
spans = span_exporter.get_finished_spans()
list_spans = [s for s in spans if s.name == "prompts/list"]
assert len(list_spans) == 1
span = list_spans[0]
assert span.kind == SpanKind.CLIENT
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "prompts/list"
# endregion
# region tools/call span
async def test_mcp_tools_call_creates_client_span_when_no_parent(span_exporter: InMemorySpanExporter):
"""Direct call_tool() without FunctionTool wrapper creates new MCP CLIENT span."""
tool = _make_connected_mcp_tool()
tool.session.call_tool = AsyncMock(return_value=_make_call_tool_result("hello"))
span_exporter.clear()
result = await tool.call_tool("get-weather", city="Seattle")
assert result is not None
spans = span_exporter.get_finished_spans()
call_spans = [s for s in spans if "tools/call" in s.name]
assert len(call_spans) == 1
span = call_spans[0]
assert span.kind == SpanKind.CLIENT
assert span.name == "tools/call get-weather"
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "tools/call"
assert span.attributes[OtelAttr.TOOL_NAME] == "get-weather"
async def test_mcp_tools_call_tool_error_sets_error_type(span_exporter: InMemorySpanExporter):
"""When CallToolResult.isError is true, error.type should be 'tool_error' per MCP spec."""
tool = _make_connected_mcp_tool()
tool.session.call_tool = AsyncMock(return_value=_make_call_tool_result("bad input", is_error=True))
span_exporter.clear()
with pytest.raises(ToolExecutionException):
await tool.call_tool("get-weather", city="invalid")
spans = span_exporter.get_finished_spans()
call_spans = [s for s in spans if "tools/call" in s.name]
assert len(call_spans) == 1
span = call_spans[0]
assert span.attributes.get(OtelAttr.ERROR_TYPE) == "tool_error"
assert span.status.status_code == StatusCode.ERROR
async def test_mcp_tools_call_mcp_error_sets_error_type(span_exporter: InMemorySpanExporter):
"""When session.call_tool() raises McpError, error.type should be the exception class name."""
tool = _make_connected_mcp_tool()
tool.session.call_tool = AsyncMock(side_effect=McpError(ErrorData(code=-32600, message="invalid request")))
span_exporter.clear()
with pytest.raises(ToolExecutionException):
await tool.call_tool("get-weather")
spans = span_exporter.get_finished_spans()
call_spans = [s for s in spans if "tools/call" in s.name]
assert len(call_spans) == 1
span = call_spans[0]
assert span.attributes.get(OtelAttr.ERROR_TYPE) == "McpError"
assert span.status.status_code == StatusCode.ERROR
# endregion
# region prompts/get span
async def test_mcp_prompts_get_creates_client_span(span_exporter: InMemorySpanExporter):
"""get_prompt() should always create a new MCP CLIENT span (not enrich execute_tool)."""
tool = _make_connected_mcp_tool()
tool.session.get_prompt = AsyncMock(return_value=_make_get_prompt_result("code analysis"))
span_exporter.clear()
result = await tool.get_prompt("analyze-code", language="python")
assert "code analysis" in result
spans = span_exporter.get_finished_spans()
prompt_spans = [s for s in spans if "prompts/get" in s.name]
assert len(prompt_spans) == 1
span = prompt_spans[0]
assert span.kind == SpanKind.CLIENT
assert span.name == "prompts/get analyze-code"
assert span.attributes[OtelAttr.MCP_METHOD_NAME] == "prompts/get"
assert span.attributes[OtelAttr.PROMPT_NAME] == "analyze-code"
async def test_mcp_prompts_get_mcp_error_sets_error_type(span_exporter: InMemorySpanExporter):
"""When session.get_prompt() raises McpError, the span should have error.type and ERROR status."""
tool = _make_connected_mcp_tool()
tool.session.get_prompt = AsyncMock(
side_effect=McpError(ErrorData(code=-32602, message="prompt not found"))
)
span_exporter.clear()
with pytest.raises(ToolExecutionException):
await tool.get_prompt("missing-prompt")
spans = span_exporter.get_finished_spans()
prompt_spans = [s for s in spans if "prompts/get" in s.name]
assert len(prompt_spans) == 1
span = prompt_spans[0]
assert span.attributes.get(OtelAttr.ERROR_TYPE) == "McpError"
assert span.status.status_code == StatusCode.ERROR
# endregion
# region transport attributes
def test_mcp_stdio_tool_transport_attributes():
"""MCPStdioTool should have network.transport='pipe'."""
tool = MCPStdioTool(name="test", command="python")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.NETWORK_TRANSPORT] == "pipe"
assert OtelAttr.ADDRESS not in attrs
def test_mcp_http_tool_transport_attributes():
"""MCPStreamableHTTPTool should have tcp transport and URL-based server address/port."""
tool = MCPStreamableHTTPTool(name="test", url="https://api.example.com:8443/mcp")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.NETWORK_TRANSPORT] == "tcp"
assert attrs[OtelAttr.NETWORK_PROTOCOL_NAME] == "http"
assert attrs[OtelAttr.ADDRESS] == "api.example.com"
assert attrs[OtelAttr.PORT] == 8443
def test_mcp_http_tool_default_port():
"""MCPStreamableHTTPTool should default to 443 for https."""
tool = MCPStreamableHTTPTool(name="test", url="https://api.example.com/mcp")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.PORT] == 443
def test_mcp_http_tool_http_default_port():
"""MCPStreamableHTTPTool should default to 80 for http."""
tool = MCPStreamableHTTPTool(name="test", url="http://localhost/mcp")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.PORT] == 80
def test_mcp_websocket_tool_transport_attributes():
"""MCPWebsocketTool should have tcp transport and URL-based server address/port."""
tool = MCPWebsocketTool(name="test", url="wss://ws.example.com:9090/mcp")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.NETWORK_TRANSPORT] == "tcp"
assert attrs[OtelAttr.NETWORK_PROTOCOL_NAME] == "websocket"
assert attrs[OtelAttr.ADDRESS] == "ws.example.com"
assert attrs[OtelAttr.PORT] == 9090
def test_mcp_websocket_tool_default_port():
"""MCPWebsocketTool should default to 443 for wss."""
tool = MCPWebsocketTool(name="test", url="wss://ws.example.com/mcp")
attrs = tool._mcp_base_span_attributes()
assert attrs[OtelAttr.PORT] == 443
# endregion
# region observability disabled
@pytest.mark.parametrize("enable_instrumentation", [False], indirect=True)
async def test_mcp_spans_not_created_when_observability_disabled(span_exporter: InMemorySpanExporter):
"""No MCP spans should be created when observability is disabled."""
tool = _make_connected_mcp_tool()
tool.session.list_tools = AsyncMock(return_value=_make_tool_list_result())
tool.session.call_tool = AsyncMock(return_value=_make_call_tool_result("ok"))
span_exporter.clear()
await tool.load_tools()
await tool.call_tool("get-weather", city="Seattle")
spans = span_exporter.get_finished_spans()
assert len(spans) == 0
# endregion
@@ -25,6 +25,7 @@ from agent_framework import (
prepend_agent_framework_to_user_agent,
tool,
)
from agent_framework._serialization import make_json_safe
from agent_framework.observability import (
ROLE_EVENT_MAP,
AgentTelemetryLayer,
@@ -3195,17 +3196,15 @@ def test_capture_messages_with_prepared_request_info_function_call_arguments(spa
from opentelemetry import trace
from agent_framework import WorkflowAgent
@dataclasses.dataclass
class HandoffRequest:
target_agent: str
reason: str
arguments = WorkflowAgent.RequestInfoFunctionArgs(
request_id="call_dc",
data=HandoffRequest(target_agent="helper", reason="overflow"),
).to_dict()
arguments = {
"request_id": "call_dc",
"data": make_json_safe(HandoffRequest(target_agent="helper", reason="overflow")),
}
msg = Message(
role="assistant",
contents=[
@@ -699,3 +699,171 @@ async def test_resolve_executor_kwargs_empty_per_executor_does_not_fallback_to_g
resolved = {"exec_a": {}, GLOBAL_KWARGS_KEY: {"global_key": "global_val"}}
result = executor._resolve_executor_kwargs(resolved) # pyright: ignore[reportPrivateUsage]
assert result == {}
# region Tool approval emission
class _ApprovalEmittingAgent(BaseAgent):
"""Agent that returns a single ``function_approval_request`` Content.
Used to verify that ``AgentExecutor`` does *not* surface the approval
payload via both an ``output`` event and a ``request_info`` event in the
same superstep only the ``request_info`` event must carry it.
"""
def __init__(
self,
*,
approval_request_id: str = "apr_1",
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
**kwargs: Any,
):
super().__init__(**kwargs)
self._approval_request_id = approval_request_id
self._tool_name = tool_name
self._tool_arguments: dict[str, Any] = tool_arguments or {"path": "/tmp/secret.txt"}
self.run_count = 0
def _build_approval_content(self) -> Content:
function_call = Content.from_function_call(
call_id=self._approval_request_id,
name=self._tool_name,
arguments=self._tool_arguments,
)
return Content.from_function_approval_request(id=self._approval_request_id, function_call=function_call)
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: AgentRunInputs | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: AgentRunInputs | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]] | ResponseStream[AgentResponseUpdate, AgentResponse[Any]]:
self.run_count += 1
approval = self._build_approval_content()
if stream:
async def _stream() -> AsyncIterable[AgentResponseUpdate]:
yield AgentResponseUpdate(contents=[approval], role="assistant")
return ResponseStream(_stream(), finalizer=AgentResponse.from_updates)
async def _run() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", [approval])])
return _run()
def _has_approval_payload(event: WorkflowEvent[Any]) -> bool:
"""Return True if the event's data carries a ``function_approval_request`` content."""
data: Any = event.data
def _contents_of(value: Any) -> list[Content]:
if isinstance(value, AgentResponseUpdate):
return list(value.contents)
if isinstance(value, AgentResponse):
return [c for m in value.messages for c in m.contents]
if isinstance(value, AgentExecutorResponse):
return [c for m in value.agent_response.messages for c in m.contents]
if isinstance(value, Message):
return list(value.contents)
if isinstance(value, Content):
return [value]
return []
return any(c.type == "function_approval_request" for c in _contents_of(data))
async def test_agent_executor_does_not_double_emit_approval_non_streaming() -> None:
"""Non-streaming: approval payload must only appear in the ``request_info`` event.
Regression test for the bug where ``AgentExecutor._run_agent`` first
``yield_output``-ed the response (carrying the approval Content) and then
additionally emitted a ``request_info`` event for the same payload.
"""
agent = _ApprovalEmittingAgent(id="approve_agent", name="ApproveAgent", approval_request_id="apr_ns_1")
executor = AgentExecutor(agent, id="approve_exec")
workflow = WorkflowBuilder(start_executor=executor).build()
request_info_events: list[WorkflowEvent[Any]] = []
output_events: list[WorkflowEvent[Any]] = []
for event in await workflow.run("please delete it"):
if event.type == "request_info":
request_info_events.append(event)
elif event.type == "output":
output_events.append(event)
assert len(request_info_events) == 1
assert _has_approval_payload(request_info_events[0])
# The approval payload must not also be surfaced as a workflow output.
assert not any(_has_approval_payload(e) for e in output_events)
assert agent.run_count == 1
async def test_agent_executor_does_not_double_emit_approval_streaming() -> None:
"""Streaming: per-update approval payload must not be ``yield_output``-ed."""
agent = _ApprovalEmittingAgent(id="approve_agent_s", name="ApproveAgentS", approval_request_id="apr_st_1")
executor = AgentExecutor(agent, id="approve_exec_s")
workflow = WorkflowBuilder(start_executor=executor).build()
request_info_events: list[WorkflowEvent[Any]] = []
output_events: list[WorkflowEvent[Any]] = []
async for event in workflow.run("please delete it", stream=True):
if event.type == "request_info":
request_info_events.append(event)
elif event.type == "output":
output_events.append(event)
assert len(request_info_events) == 1
assert _has_approval_payload(request_info_events[0])
assert not any(_has_approval_payload(e) for e in output_events)
assert agent.run_count == 1
async def test_agent_executor_request_info_uses_user_input_request_id() -> None:
"""``ctx.request_info`` must register the request under the agent's approval id.
This makes the workflow's pending-request id round-trip with the
``function_approval_response.id`` the caller echoes back, so
``Workflow._send_responses_internal`` can look it up directly.
"""
agent = _ApprovalEmittingAgent(id="approve_agent_id", name="ApproveAgentId", approval_request_id="apr_match")
executor = AgentExecutor(agent, id="approve_exec_id")
workflow = WorkflowBuilder(start_executor=executor).build()
request_info_events: list[WorkflowEvent[Any]] = []
async for event in workflow.run("please delete it", stream=True):
if event.type == "request_info":
request_info_events.append(event)
assert len(request_info_events) == 1
assert request_info_events[0].request_id == "apr_match"
# endregion Tool approval emission
@@ -1,6 +1,5 @@
# Copyright (c) Microsoft. All rights reserved.
import json
import uuid
from collections.abc import Awaitable, Sequence
from dataclasses import dataclass
@@ -30,6 +29,20 @@ from agent_framework import (
handler,
response_handler,
)
from agent_framework._workflows._typing_utils import deserialize_type
@dataclass
class HandoffRequest:
"""Module-level dataclass used by request_info tests.
Defined at module scope (not nested inside a test method) so
``serialize_type``/``deserialize_type`` can round-trip the request_type via
the importable qualified name ``tests.workflow.test_workflow_agent.HandoffRequest``.
"""
target_agent: str
reason: str
class SimpleExecutor(Executor):
@@ -240,52 +253,45 @@ class TestWorkflowAgent:
# Should have received an approval request for the request info
assert len(updates) > 0
approval_update: AgentResponseUpdate | None = None
request_update: AgentResponseUpdate | None = None
for update in updates:
if any(content.type == "function_approval_request" for content in update.contents):
approval_update = update
if any(content.type == "function_call" for content in update.contents):
request_update = update
break
assert approval_update is not None, "Should have received a request_info approval request"
assert request_update is not None, "Should have received a request_info wrapped in a function_call content"
function_call = next(content for content in approval_update.contents if content.type == "function_call")
approval_request = next(
content for content in approval_update.contents if content.type == "function_approval_request"
)
request_function_call = next(content for content in request_update.contents if content.type == "function_call")
assert request_function_call.call_id is not None
# Verify the function call has expected structure
assert function_call.call_id is not None
assert function_call.name == "request_info"
assert isinstance(function_call.arguments, dict)
assert function_call.arguments.get("request_id") == approval_request.id
assert request_function_call.name == WorkflowAgent.REQUEST_INFO_FUNCTION_NAME
assert isinstance(request_function_call.arguments, dict)
assert request_function_call.arguments.get("request_id") is not None
assert request_function_call.arguments.get("request_event") is not None
request_event = request_function_call.arguments["request_event"]
assert request_event.get("type") == "request_info"
assert deserialize_type(request_event.get("response_type")) is str
# Approval request should reference the same function call
assert approval_request.id is not None
assert approval_request.function_call is not None
assert approval_request.function_call.call_id == function_call.call_id
assert approval_request.function_call.name == function_call.name
deserialized_args = WorkflowAgent.RequestInfoFunctionArgs.from_dict(request_function_call.arguments)
assert deserialized_args.request_id == request_function_call.call_id
assert isinstance(deserialized_args.request_event, WorkflowEvent)
assert deserialized_args.request_event.type == "request_info"
assert deserialized_args.request_event.data == "Mock request data"
assert deserialized_args.request_event.response_type is str
# Verify the request is tracked in pending_requests
assert len(agent.pending_requests) == 1
assert function_call.call_id in agent.pending_requests
pending_requests = await workflow._runner_context.get_pending_request_info_events()
assert len(pending_requests) == 1
assert request_function_call.call_id in pending_requests
# Now provide an approval response with updated arguments to test continuation
response_args = WorkflowAgent.RequestInfoFunctionArgs(
request_id=approval_request.id,
data="User provided answer",
).to_dict()
approval_response = Content.from_function_approval_response(
approved=True,
id=approval_request.id,
function_call=Content.from_function_call(
call_id=function_call.call_id,
name=function_call.name,
arguments=response_args,
),
# Now provide a function result response with updated arguments to test continuation
function_result = Content.from_function_result(
call_id=request_function_call.call_id,
result="Mock response to request info",
)
response_message = Message(role="user", contents=[approval_response])
response_message = Message(role="user", contents=[function_result])
# Continue the workflow with the response
continuation_result = await agent.run(response_message)
@@ -294,16 +300,11 @@ class TestWorkflowAgent:
assert isinstance(continuation_result, AgentResponse)
# Verify cleanup - pending requests should be cleared after function response handling
assert len(agent.pending_requests) == 0
pending_requests = await workflow._runner_context.get_pending_request_info_events()
assert len(pending_requests) == 0
def test_request_info_dataclass_arguments_are_serialized_when_content_is_created(self) -> None:
"""Test WorkflowAgent prepares request_info arguments before observability captures messages."""
@dataclass
class HandoffRequest:
target_agent: str
reason: str
executor = SimpleExecutor(id="executor1", response_text="Response")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="Request Test Agent")
@@ -314,14 +315,367 @@ class TestWorkflowAgent:
response_type=str,
)
function_call, approval_request = agent._process_request_info_event(event) # pyright: ignore[reportPrivateUsage]
request_function_call = agent._process_request_info_event(event) # pyright: ignore[reportPrivateUsage]
assert function_call.arguments == {
"request_id": "request_123",
"data": {"target_agent": "helper", "reason": "overflow"},
}
assert approval_request.function_call is function_call
assert json.loads(json.dumps(function_call.arguments)) == function_call.arguments
assert request_function_call.call_id == "request_123"
assert isinstance(request_function_call.arguments, dict)
assert request_function_call.arguments.get("request_event") is not None
request_event = request_function_call.arguments["request_event"]
assert request_event.get("type") == "request_info"
assert request_event.get("request_id") == "request_123"
assert request_event.get("source_executor_id") == "executor1"
assert deserialize_type(request_event.get("response_type")) is str
assert request_event.get("data") == HandoffRequest(target_agent="helper", reason="overflow")
deserialized_args = WorkflowAgent.RequestInfoFunctionArgs.from_dict(request_function_call.arguments)
assert deserialized_args.request_id == "request_123"
assert isinstance(deserialized_args.request_event, WorkflowEvent)
assert deserialized_args.request_event.type == "request_info"
assert deserialized_args.request_event.data == HandoffRequest(target_agent="helper", reason="overflow")
assert deserialized_args.request_event.response_type is str
def test_process_request_info_event_passes_through_function_approval_request(self) -> None:
"""If the event data is already a function approval request, it is forwarded unchanged.
Tool-approval requests emitted by an inner agent surface as ``Content``
objects with ``user_input_request=True``. ``WorkflowAgent`` must not
re-wrap these inside a synthesized ``request_info`` function call;
instead it should return the original content as-is so callers can
respond with a matching ``function_approval_response``.
"""
executor = SimpleExecutor(id="executor1", response_text="Response")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Passthrough Agent")
approval_id = "approval-passthrough-1"
inner_function_call = Content.from_function_call(
call_id="tool-call-1",
name="delete_file",
arguments={"path": "/tmp/x"},
)
approval_request = Content.from_function_approval_request(
id=approval_id,
function_call=inner_function_call,
)
event = WorkflowEvent.request_info(
request_id=approval_id,
source_executor_id="executor1",
request_data=approval_request,
response_type=Content,
)
result = agent._process_request_info_event(event) # pyright: ignore[reportPrivateUsage]
# The original FunctionApprovalRequestContent is returned as-is — same
# instance, with the original tool name preserved (NOT replaced by the
# synthesized REQUEST_INFO_FUNCTION_NAME).
assert result is approval_request
assert result.type == "function_approval_request"
assert result.id == approval_id
assert result.user_input_request is True
assert result.function_call is inner_function_call # type: ignore[attr-defined]
assert result.function_call.name == "delete_file" # type: ignore[attr-defined]
assert result.function_call.name != WorkflowAgent.REQUEST_INFO_FUNCTION_NAME # type: ignore[attr-defined]
def test_extract_function_responses_passes_through_approval_response_approved(self) -> None:
"""A function_approval_response with approved=True is keyed by content.id and forwarded as-is.
After the refactor, ``WorkflowAgent`` no longer unwraps a synthesized
``request_info`` function call from approval responses the response
content is routed straight back to the workflow under its own ``id``,
which matches the pending request id surfaced by
``_process_request_info_event``.
"""
executor = SimpleExecutor(id="executor1", response_text="Response")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Response Agent")
approval_id = "approval-response-approved-1"
inner_function_call = Content.from_function_call(
call_id="tool-call-1",
name="delete_file",
arguments={"path": "/tmp/x"},
)
approval_request = Content.from_function_approval_request(
id=approval_id,
function_call=inner_function_call,
)
approval_response = approval_request.to_function_approval_response(approved=True) # type: ignore[attr-defined]
message = Message(role="user", contents=[approval_response])
responses = agent._extract_function_responses([message]) # pyright: ignore[reportPrivateUsage]
assert set(responses.keys()) == {approval_id}
assert responses[approval_id] is approval_response
assert responses[approval_id].approved is True # type: ignore[attr-defined]
def test_extract_function_responses_passes_through_approval_response_denied(self) -> None:
"""A function_approval_response with approved=False is forwarded the same way as an approval.
Only the ``approved`` flag changes routing back to the workflow is
identical for accept and reject paths.
"""
executor = SimpleExecutor(id="executor1", response_text="Response")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Response Agent")
approval_id = "approval-response-denied-1"
inner_function_call = Content.from_function_call(
call_id="tool-call-2",
name="send_email",
arguments={"to": "alice@example.com"},
)
approval_request = Content.from_function_approval_request(
id=approval_id,
function_call=inner_function_call,
)
approval_response = approval_request.to_function_approval_response(approved=False) # type: ignore[attr-defined]
message = Message(role="user", contents=[approval_response])
responses = agent._extract_function_responses([message]) # pyright: ignore[reportPrivateUsage]
assert set(responses.keys()) == {approval_id}
assert responses[approval_id] is approval_response
assert responses[approval_id].approved is False # type: ignore[attr-defined]
async def test_function_approval_request_flows_end_to_end_approved(self) -> None:
"""End-to-end: an executor emits a function_approval_request, the agent
forwards it unchanged, and an ``approved=True`` response resumes the workflow.
This exercises the full pass-through path:
``ctx.request_info(approval_content, ...)`` -> ``WorkflowAgent`` surfaces
the original ``FunctionApprovalRequestContent`` -> caller responds with a
``FunctionApprovalResponseContent`` -> ``WorkflowAgent`` routes it back
to the workflow which delivers it to the executor's ``@response_handler``.
"""
approval_id = "e2e-approval-1"
inner_function_call = Content.from_function_call(
call_id="tool-call-e2e-1",
name="delete_file",
arguments={"path": "/tmp/x"},
)
approval_request = Content.from_function_approval_request(
id=approval_id,
function_call=inner_function_call,
)
class ApprovalRequestingExecutor(Executor):
@handler
async def handle_message(self, _: list[Message], ctx: WorkflowContext) -> None:
await ctx.request_info(approval_request, Content, request_id=approval_id)
@response_handler
async def handle_response(
self,
original_request: Content,
response: Content,
ctx: WorkflowContext[Never, AgentResponse],
) -> None:
assert response.type == "function_approval_response"
assert response.id == approval_id # type: ignore[attr-defined]
approved = bool(response.approved) # type: ignore[attr-defined]
tool_name = original_request.function_call.name # type: ignore[attr-defined]
await ctx.yield_output(
AgentResponse(
messages=[
Message(
role="assistant",
contents=[Content.from_text(text=f"{tool_name} approved={approved}")],
)
]
)
)
executor = ApprovalRequestingExecutor(id="approval_requester")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="E2E Approval Agent")
# First run: workflow pauses with the approval request.
first = await agent.run("please delete it")
assert isinstance(first, AgentResponse)
forwarded = next(
(
c
for m in first.messages
for c in m.contents
if c.type == "function_approval_request" and c.id == approval_id
),
None,
)
assert forwarded is approval_request, "Approval request must surface unchanged"
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id in pending
# Respond with approved=True.
approval_response = approval_request.to_function_approval_response(approved=True) # type: ignore[attr-defined]
final = await agent.run(Message(role="user", contents=[approval_response]))
assert isinstance(final, AgentResponse)
final_text = " ".join(m.text or "" for m in final.messages)
assert "delete_file approved=True" in final_text
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id not in pending
async def test_function_approval_request_flows_end_to_end_denied(self) -> None:
"""End-to-end denied path: ``approved=False`` is delivered to the executor's
response handler so the workflow can branch on the rejection."""
approval_id = "e2e-approval-deny-1"
inner_function_call = Content.from_function_call(
call_id="tool-call-e2e-deny-1",
name="send_email",
arguments={"to": "alice@example.com"},
)
approval_request = Content.from_function_approval_request(
id=approval_id,
function_call=inner_function_call,
)
class ApprovalRequestingExecutor(Executor):
@handler
async def handle_message(self, _: list[Message], ctx: WorkflowContext) -> None:
await ctx.request_info(approval_request, Content, request_id=approval_id)
@response_handler
async def handle_response(
self,
original_request: Content,
response: Content,
ctx: WorkflowContext[Never, AgentResponse],
) -> None:
assert response.type == "function_approval_response"
assert response.id == approval_id # type: ignore[attr-defined]
approved = bool(response.approved) # type: ignore[attr-defined]
tool_name = original_request.function_call.name # type: ignore[attr-defined]
await ctx.yield_output(
AgentResponse(
messages=[
Message(
role="assistant",
contents=[Content.from_text(text=f"{tool_name} approved={approved}")],
)
]
)
)
executor = ApprovalRequestingExecutor(id="approval_requester_deny")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="E2E Approval Deny Agent")
first = await agent.run("please send")
assert isinstance(first, AgentResponse)
forwarded = next(
(
c
for m in first.messages
for c in m.contents
if c.type == "function_approval_request" and c.id == approval_id
),
None,
)
assert forwarded is approval_request
# Respond with approved=False.
approval_response = approval_request.to_function_approval_response(approved=False) # type: ignore[attr-defined]
final = await agent.run(Message(role="user", contents=[approval_response]))
assert isinstance(final, AgentResponse)
final_text = " ".join(m.text or "" for m in final.messages)
assert "send_email approved=False" in final_text
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id not in pending
async def test_request_info_non_approval_flows_end_to_end(self) -> None:
"""End-to-end: when request data is not a function approval content, the
agent surfaces a synthesized ``function_call`` (name=REQUEST_INFO_FUNCTION_NAME)
and routes a matching ``function_result`` back to the executor.
"""
captured: dict[str, Any] = {}
class HandoffRequestingExecutor(Executor):
@handler
async def handle_message(self, _: list[Message], ctx: WorkflowContext) -> None:
await ctx.request_info(
HandoffRequest(target_agent="helper", reason="overflow"),
str,
)
@response_handler
async def handle_response(
self,
original_request: HandoffRequest,
response: str,
ctx: WorkflowContext[Never, AgentResponse],
) -> None:
captured["original"] = original_request
captured["response"] = response
await ctx.yield_output(
AgentResponse(
messages=[
Message(
role="assistant",
contents=[
Content.from_text(text=f"handoff to {original_request.target_agent}: {response}")
],
)
]
)
)
executor = HandoffRequestingExecutor(id="handoff_requester")
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="E2E Handoff Agent")
# First run: workflow pauses with a synthesized request_info function_call.
first = await agent.run("start handoff")
assert isinstance(first, AgentResponse)
function_call = next(
(
c
for m in first.messages
for c in m.contents
if c.type == "function_call" and c.name == WorkflowAgent.REQUEST_INFO_FUNCTION_NAME
),
None,
)
assert function_call is not None, "Expected a synthesized request_info function_call"
assert function_call.call_id is not None
assert isinstance(function_call.arguments, dict)
request_id = function_call.arguments["request_id"]
assert function_call.call_id == request_id
request_payload = function_call.arguments["request_event"]
assert request_payload.get("type") == "request_info"
assert request_payload.get("data") == HandoffRequest(target_agent="helper", reason="overflow")
deserialized_args = WorkflowAgent.RequestInfoFunctionArgs.from_dict(function_call.arguments)
assert deserialized_args.request_id == request_id
assert isinstance(deserialized_args.request_event, WorkflowEvent)
assert deserialized_args.request_event.type == "request_info"
assert deserialized_args.request_event.data == HandoffRequest(target_agent="helper", reason="overflow")
assert deserialized_args.request_event.response_type is str
pending = await workflow._runner_context.get_pending_request_info_events()
assert request_id in pending
# Respond with a function_result keyed by the call_id.
function_result = Content.from_function_result(call_id=request_id, result="ok-do-it")
final = await agent.run(Message(role="user", contents=[function_result]))
assert isinstance(final, AgentResponse)
final_text = " ".join(m.text or "" for m in final.messages)
assert "handoff to helper: ok-do-it" in final_text
# The executor's response handler received the original request and the response.
assert isinstance(captured.get("original"), HandoffRequest)
assert captured["original"].target_agent == "helper"
assert captured["response"] == "ok-do-it"
pending = await workflow._runner_context.get_pending_request_info_events()
assert request_id not in pending
def test_workflow_as_agent_method(self) -> None:
"""Test that Workflow.as_agent() creates a properly configured WorkflowAgent."""
@@ -1592,3 +1946,406 @@ class TestWorkflowAgentMergeUpdates:
# Order: text (user), text (assistant), function_result (orphan at end)
assert content_types == ["text", "text", "function_result"]
class _ToolApprovalMockAgent(SupportsAgentRun):
"""Mock agent whose first run returns a FunctionApprovalRequestContent.
Subsequent runs (after receiving an approval response in the input messages)
return a final assistant text response that echoes the approved arguments.
This mirrors a real agent whose tool invocation requires user approval.
"""
def __init__(
self,
name: str,
*,
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
approval_request_ids: Sequence[str] | None = None,
) -> None:
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._tool_name = tool_name
self._tool_arguments = tool_arguments or {"path": "/tmp/example"}
# Pre-allocated request ids so the test can verify what the WorkflowAgent forwards.
self._approval_request_ids: list[str] = list(approval_request_ids) if approval_request_ids else []
self.run_count = 0
# Inputs received on the most recent (continuation) run, for assertions.
self.last_run_messages: list[Message] = []
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
return AgentSession()
def _next_request_id(self) -> str:
if self._approval_request_ids:
return self._approval_request_ids.pop(0)
return str(uuid.uuid4())
def _build_approval_request(self) -> Content:
request_id = self._next_request_id()
function_call = Content.from_function_call(
call_id=request_id,
name=self._tool_name,
arguments=self._tool_arguments,
)
return Content.from_function_approval_request(id=request_id, function_call=function_call)
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
return self._run_stream(messages=messages, session=session, **kwargs)
return self._run(messages=messages, session=session, **kwargs)
def _normalize(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None,
) -> list[Message]:
if messages is None:
return []
if isinstance(messages, str):
return [Message(role="user", contents=[Content.from_text(text=messages)])]
if isinstance(messages, Message):
return [messages]
if isinstance(messages, Content):
return [Message(role="user", contents=[messages])]
result: list[Message] = []
for item in messages:
if isinstance(item, Message):
result.append(item)
elif isinstance(item, Content):
result.append(Message(role="user", contents=[item]))
else:
result.append(Message(role="user", contents=[Content.from_text(text=item)]))
return result
def _approval_responses_in(self, messages: list[Message]) -> list[Content]:
approvals: list[Content] = []
for msg in messages:
for content in msg.contents:
if content.type == "function_approval_response":
approvals.append(content)
return approvals
async def _run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
session: AgentSession | None = None,
**kwargs: Any,
) -> AgentResponse:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
approvals = self._approval_responses_in(normalized)
if approvals:
# Continuation: reflect approved arguments in the final response text.
approved_text = "; ".join(
f"approved={a.approved} id={a.id}" # type: ignore[attr-defined]
for a in approvals
)
return AgentResponse(messages=[Message("assistant", [Content.from_text(text=f"done ({approved_text})")])])
# First run: ask for tool approval.
approval = self._build_approval_request()
return AgentResponse(messages=[Message("assistant", [approval])])
def _run_stream(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
session: AgentSession | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
approvals = self._approval_responses_in(normalized)
async def _iter():
if approvals:
approved_text = "; ".join(
f"approved={a.approved} id={a.id}" # type: ignore[attr-defined]
for a in approvals
)
yield AgentResponseUpdate(
contents=[Content.from_text(text=f"done ({approved_text})")],
role="assistant",
author_name=self.name,
)
return
approval = self._build_approval_request()
yield AgentResponseUpdate(
contents=[approval],
role="assistant",
author_name=self.name,
)
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
class TestWorkflowAgentToolApproval:
"""Tests for tool-approval requests bubbling through WorkflowAgent.
Covers the case where a workflow contains an AgentExecutor whose underlying
agent emits a FunctionApprovalRequestContent (tool needing user approval).
The WorkflowAgent must:
* forward the original FunctionApprovalRequestContent unchanged (no
wrapping inside a synthesized 'request_info' function call), and
* route a subsequent FunctionApprovalResponseContent back to the
AgentExecutor so the agent can resume.
"""
def _find_approval_request(
self,
contents: Sequence[Content],
tool_name: str,
) -> Content | None:
for content in contents:
if (
content.type == "function_approval_request"
and getattr(content.function_call, "name", None) == tool_name # type: ignore[attr-defined]
):
return content
return None
async def test_tool_approval_request_forwarded_unchanged(self) -> None:
"""The agent's FunctionApprovalRequestContent surfaces verbatim (not re-wrapped)."""
approval_id = "approval-abc-123"
mock_agent = _ToolApprovalMockAgent(
name="approval-agent",
tool_name="delete_file",
tool_arguments={"path": "/tmp/secret.txt"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Test Agent")
result = await agent.run("please delete the file")
assert isinstance(result, AgentResponse)
# Locate the approval request emitted by the WorkflowAgent.
all_contents: list[Content] = [c for m in result.messages for c in m.contents]
approval = self._find_approval_request(all_contents, tool_name="delete_file")
assert approval is not None, "WorkflowAgent did not forward the tool approval request"
# The id and inner function_call must match what the underlying agent produced
# — i.e. the WorkflowAgent must NOT have re-wrapped it inside a synthesized
# 'request_info' approval request.
assert approval.id == approval_id
function_call = approval.function_call # type: ignore[attr-defined]
assert function_call is not None
assert function_call.name == "delete_file"
assert function_call.name != WorkflowAgent.REQUEST_INFO_FUNCTION_NAME
assert function_call.arguments == {"path": "/tmp/secret.txt"}
# The agent must be paused awaiting the approval response.
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id in pending
async def test_tool_approval_request_forwarded_unchanged_streaming(self) -> None:
"""Streaming variant: the approval request is forwarded as-is in updates."""
approval_id = "approval-stream-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-agent-stream",
tool_name="send_email",
tool_arguments={"to": "alice@example.com"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Stream Agent")
updates: list[AgentResponseUpdate] = []
async for update in agent.run("hi", stream=True):
updates.append(update)
approval_updates = [u for u in updates if any(c.type == "function_approval_request" for c in u.contents)]
assert approval_updates, "Streaming did not surface a tool approval request"
approval = self._find_approval_request(approval_updates[-1].contents, tool_name="send_email")
assert approval is not None
assert approval.id == approval_id
function_call = approval.function_call # type: ignore[attr-defined]
assert function_call is not None
assert function_call.name == "send_email"
assert function_call.name != WorkflowAgent.REQUEST_INFO_FUNCTION_NAME
assert function_call.arguments == {"to": "alice@example.com"}
async def test_tool_approval_response_resumes_agent(self) -> None:
"""Sending the approval response back resumes the agent and clears pending requests."""
approval_id = "approval-resume-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-resume-agent",
tool_name="delete_file",
tool_arguments={"path": "/tmp/x"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Resume Agent")
first_result = await agent.run("delete it")
approval = self._find_approval_request(
[c for m in first_result.messages for c in m.contents],
tool_name="delete_file",
)
assert approval is not None
assert mock_agent.run_count == 1
# Build the approval response. NOTE: the inner function_call's name is the
# original tool name ('delete_file'), NOT 'request_info'. This exercises the
# branch in WorkflowAgent._extract_function_responses that routes raw
# tool-approval responses straight through using content.id.
approval_response = approval.to_function_approval_response(approved=True) # type: ignore[attr-defined]
response_message = Message(role="user", contents=[approval_response])
final_result = await agent.run(response_message)
assert isinstance(final_result, AgentResponse)
# The mock agent should have been invoked a second time and seen the
# approval response in its inputs.
assert mock_agent.run_count == 2
approvals_seen = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approvals_seen) == 1
assert approvals_seen[0].id == approval_id # type: ignore[attr-defined]
assert approvals_seen[0].approved is True # type: ignore[attr-defined]
# The pending approval should now be cleared.
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id not in pending
# The final assistant message reflects the resumption.
final_text = " ".join(m.text or "" for m in final_result.messages)
assert "done" in final_text
assert approval_id in final_text
async def test_tool_approval_response_rejected_resumes_agent(self) -> None:
"""Rejection path: ``approved=False`` is forwarded to the inner agent and clears the pending request.
The WorkflowAgent must route a rejection response back to the paused
``AgentExecutor`` exactly the same way as an approval only the
``approved`` flag differs. The inner agent decides what to do with it.
"""
approval_id = "approval-reject-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-reject-agent",
tool_name="delete_file",
tool_arguments={"path": "/tmp/x"},
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Reject Agent")
first_result = await agent.run("delete it")
approval = self._find_approval_request(
[c for m in first_result.messages for c in m.contents],
tool_name="delete_file",
)
assert approval is not None
assert mock_agent.run_count == 1
# Reject the tool invocation.
approval_response = approval.to_function_approval_response(approved=False) # type: ignore[attr-defined]
response_message = Message(role="user", contents=[approval_response])
final_result = await agent.run(response_message)
assert isinstance(final_result, AgentResponse)
# The inner agent must have been resumed and seen ``approved=False``.
assert mock_agent.run_count == 2
approvals_seen = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approvals_seen) == 1
assert approvals_seen[0].id == approval_id # type: ignore[attr-defined]
assert approvals_seen[0].approved is False # type: ignore[attr-defined]
# Pending approval cleared regardless of approve/reject.
pending = await workflow._runner_context.get_pending_request_info_events()
assert approval_id not in pending
# The final assistant message reflects the rejection.
final_text = " ".join(m.text or "" for m in final_result.messages)
assert "approved=False" in final_text
assert approval_id in final_text
async def test_tool_approval_request_id_matches_pending_request(self) -> None:
"""The approval request id surfaced by WorkflowAgent matches the workflow's pending request id.
This guards the AgentExecutor change that forwards
request_id=user_input_request.id to ctx.request_info(...), which is what
allows the response routed back via WorkflowAgent to resolve the pending
request without an id-mismatch error.
"""
approval_id = "approval-id-match-1"
mock_agent = _ToolApprovalMockAgent(
name="approval-id-match-agent",
approval_request_ids=[approval_id],
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
agent = WorkflowAgent(workflow=workflow, name="Approval Id Agent")
await agent.run("go")
pending = await workflow._runner_context.get_pending_request_info_events()
# The agent's approval id is used as the workflow's pending request id.
assert list(pending.keys()) == [approval_id]
@@ -0,0 +1,149 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for the ``Workflow.status`` property."""
from dataclasses import dataclass
import pytest
from agent_framework import (
Executor,
Workflow,
WorkflowBuilder,
WorkflowContext,
WorkflowEvent,
WorkflowRunState,
handler,
response_handler,
)
from agent_framework._workflows._executor import Executor as _Executor
from agent_framework._workflows._request_info_mixin import RequestInfoMixin
class PassThroughExecutor(Executor):
"""Executor that yields its input as a workflow output and stops."""
@handler
async def passthrough(self, msg: str, ctx: WorkflowContext[str, str]) -> None:
await ctx.yield_output(msg)
class FailingExecutor(Executor):
"""Executor that raises at runtime to drive the FAILED status."""
@handler
async def fail(self, msg: int, ctx: WorkflowContext) -> None: # pragma: no cover - invoked via workflow
raise RuntimeError("boom")
@dataclass
class _ApprovalRequest:
prompt: str
request_id: str = ""
def __post_init__(self) -> None:
if not self.request_id:
import uuid
self.request_id = str(uuid.uuid4())
class ApprovalExecutor(_Executor, RequestInfoMixin):
"""Executor that issues a single request_info call and finalizes on response."""
def __init__(self, id: str = "approval"):
super().__init__(id=id)
@handler
async def start(self, message: str, ctx: WorkflowContext[str, str]) -> None:
await ctx.request_info(_ApprovalRequest(prompt=message), bool)
@response_handler
async def on_response(
self, original_request: _ApprovalRequest, approved: bool, ctx: WorkflowContext[str, str]
) -> None:
await ctx.yield_output(f"approved={approved}")
def _build_passthrough_workflow() -> Workflow:
executor = PassThroughExecutor(id="p")
return WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
def _build_failing_workflow() -> Workflow:
# FailingExecutor has no workflow_output_types, so we leave designation
# implicit; the deprecation warning is filtered at call sites that need it.
return WorkflowBuilder(start_executor=FailingExecutor(id="f")).build()
def _build_approval_workflow() -> Workflow:
executor = ApprovalExecutor(id="approval")
return WorkflowBuilder(start_executor=executor, output_from=[executor]).build()
async def test_status_default_is_idle_before_first_run():
wf = _build_passthrough_workflow()
assert wf.status is WorkflowRunState.IDLE
async def test_status_is_idle_after_successful_run():
wf = _build_passthrough_workflow()
await wf.run("hello")
assert wf.status is WorkflowRunState.IDLE
async def test_status_is_failed_after_failure():
wf = _build_failing_workflow()
with pytest.raises(RuntimeError, match="boom"):
await wf.run(0)
assert wf.status is WorkflowRunState.FAILED
async def test_status_transitions_during_streaming_run():
"""Workflow.status mirrors the most recent emitted status event."""
wf = _build_passthrough_workflow()
observed: list[WorkflowRunState] = []
async for event in wf.run("hi", stream=True):
if isinstance(event, WorkflowEvent) and event.type == "status":
# By the time a status event surfaces to the consumer, the property
# must already reflect that state (updated in lockstep with emission).
assert wf.status == event.state
observed.append(event.state) # type: ignore
# IN_PROGRESS must precede IDLE; both must appear.
assert WorkflowRunState.IN_PROGRESS in observed
assert observed[-1] is WorkflowRunState.IDLE
assert wf.status is WorkflowRunState.IDLE
async def test_status_idle_with_pending_requests_then_resolves_to_idle():
wf = _build_approval_workflow()
request_event: WorkflowEvent | None = None
async for event in wf.run("please approve", stream=True):
if isinstance(event, WorkflowEvent) and event.type == "request_info":
request_event = event
assert request_event is not None
assert wf.status is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
async for _ in wf.run(stream=True, responses={request_event.request_id: True}):
pass
assert wf.status is WorkflowRunState.IDLE
async def test_status_in_progress_pending_requests_observed_mid_run():
"""While streaming, status reaches IN_PROGRESS_PENDING_REQUESTS after a request_info event."""
wf = _build_approval_workflow()
seen_states: list[WorkflowRunState] = []
async for event in wf.run("please approve", stream=True):
if isinstance(event, WorkflowEvent) and event.type == "status":
seen_states.append(event.state) # type: ignore
assert WorkflowRunState.IN_PROGRESS in seen_states
assert WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS in seen_states
assert seen_states[-1] is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
assert wf.status is WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
+1 -1
View File
@@ -26,7 +26,7 @@ dependencies = [
"agent-framework-core>=1.8.0,<2",
"agent-framework-openai>=1.8.0,<2",
"azure-ai-inference>=1.0.0b9,<1.0.0b10",
"azure-ai-projects>=2.1.0,<3.0",
"azure-ai-projects>=2.2.0,<3.0",
]
[tool.uv]
@@ -567,7 +567,7 @@ class ResponsesHostServer(ResponsesAgentServerHost):
by the hosting infrastructure or files will be preserved upon deactivation.
"""
input_items = await context.get_input_items()
input_messages = await _items_to_messages(input_items)
input_messages = await _items_to_messages(input_items, approval_storage=self._approval_storage)
is_streaming_request = request.stream is not None and request.stream is True
_, are_options_set = _to_chat_options(request)
@@ -664,7 +664,11 @@ class ResponsesHostServer(ResponsesAgentServerHost):
checkpoint_storage=write_storage,
)
async for item in _to_outputs_for_messages(response_event_stream, response.messages):
async for item in _to_outputs_for_messages(
response_event_stream,
response.messages,
approval_storage=self._approval_storage,
):
yield item
await self._delete_not_latest_checkpoints(write_storage, self._agent.workflow.name)
@@ -685,7 +689,9 @@ class ResponsesHostServer(ResponsesAgentServerHost):
for event in tracker.handle(content):
yield event
if tracker.needs_async:
async for item in _to_outputs(response_event_stream, content):
async for item in _to_outputs(
response_event_stream, content, approval_storage=self._approval_storage
):
yield item
tracker.needs_async = False
@@ -11,24 +11,33 @@ the registered _handle_create handler.
from __future__ import annotations
import json
from collections.abc import AsyncIterator, Callable
import uuid
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
from dataclasses import dataclass
from typing import Literal, overload
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
import pytest
from agent_framework import (
AgentExecutorRequest,
AgentResponse,
AgentResponseUpdate,
AgentSession,
Content,
FileCheckpointStorage,
HistoryProvider,
Message,
RawAgent,
ResponseStream,
SupportsAgentRun,
WorkflowAgent,
WorkflowBuilder,
WorkflowCheckpoint,
WorkflowCheckpointException,
WorkflowContext,
WorkflowMessage,
executor,
)
from azure.ai.agentserver.responses import InMemoryResponseProvider
from mcp import McpError
@@ -102,7 +111,7 @@ def _make_agent(
return agent
def _make_server(agent: MagicMock, **kwargs: Any) -> ResponsesHostServer:
def _make_server(agent: Any, **kwargs: Any) -> ResponsesHostServer:
"""Create a ResponsesHostServer with an in-memory store."""
return ResponsesHostServer(agent, store=InMemoryResponseProvider(), **kwargs)
@@ -3469,3 +3478,498 @@ class TestOAuthConsentSurfacing:
# endregion
# region Workflow agent hosting (end-to-end)
class _ToolApprovalWorkflowAgentMock(SupportsAgentRun):
"""Inner agent for a hosted ``WorkflowAgent`` whose first run emits a
``FunctionApprovalRequestContent`` and whose follow-up run (after
receiving a ``FunctionApprovalResponseContent`` in its inputs) returns a
final assistant text response.
Mirrors a real agent whose tool invocation requires user approval. Used
here to exercise the full HTTP pipeline through ``ResponsesHostServer``
when the hosted agent is a ``WorkflowAgent`` containing a tool-approval
flow.
"""
def __init__(
self,
name: str,
*,
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
approval_request_ids: Sequence[str] | None = None,
final_text: str = "done",
) -> None:
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._tool_name = tool_name
self._tool_arguments = tool_arguments or {"path": "/tmp/example"}
self._approval_request_ids: list[str] = list(approval_request_ids) if approval_request_ids else []
self._final_text = final_text
self.run_count = 0
self.last_run_messages: list[Message] = []
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
return AgentSession()
def _next_request_id(self) -> str:
# Stable across calls: when the workflow checkpoint round-trips through
# restore, ``AgentExecutor`` re-invokes the inner agent during replay.
# We must surface the *same* approval request id on each invocation so
# the workflow's pending-request id matches the id the test echoes
# back as ``mcp_approval_response``.
if self._approval_request_ids:
return self._approval_request_ids[0]
return str(uuid.uuid4())
def _build_approval_request(self) -> Content:
request_id = self._next_request_id()
function_call = Content.from_function_call(
call_id=request_id,
name=self._tool_name,
arguments=self._tool_arguments,
additional_properties={"server_label": "test_server"},
)
return Content.from_function_approval_request(id=request_id, function_call=function_call)
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
if stream:
return self._run_stream(messages=messages, **kwargs)
return self._run(messages=messages, **kwargs)
@staticmethod
def _normalize(
messages: str | Content | Message | Sequence[str | Content | Message] | None,
) -> list[Message]:
if messages is None:
return []
if isinstance(messages, str):
return [Message(role="user", contents=[Content.from_text(text=messages)])]
if isinstance(messages, Message):
return [messages]
if isinstance(messages, Content):
return [Message(role="user", contents=[messages])]
result: list[Message] = []
for item in messages:
if isinstance(item, Message):
result.append(item)
elif isinstance(item, Content):
result.append(Message(role="user", contents=[item]))
else:
result.append(Message(role="user", contents=[Content.from_text(text=item)]))
return result
@staticmethod
def _approval_responses_in(messages: list[Message]) -> list[Content]:
return [c for m in messages for c in m.contents if c.type == "function_approval_response"]
async def _run(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
**kwargs: Any,
) -> AgentResponse:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
if self._approval_responses_in(normalized):
return AgentResponse(messages=[Message("assistant", [Content.from_text(text=self._final_text)])])
approval = self._build_approval_request()
return AgentResponse(messages=[Message("assistant", [approval])])
def _run_stream(
self,
messages: str | Content | Message | Sequence[str | Content | Message] | None = None,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse]:
normalized = self._normalize(messages)
self.last_run_messages = normalized
self.run_count += 1
approvals = self._approval_responses_in(normalized)
async def _iter() -> AsyncIterator[AgentResponseUpdate]:
if approvals:
yield AgentResponseUpdate(
contents=[Content.from_text(text=self._final_text)],
role="assistant",
author_name=self.name,
)
return
yield AgentResponseUpdate(
contents=[self._build_approval_request()],
role="assistant",
author_name=self.name,
)
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
def _build_text_workflow_agent(text: str) -> WorkflowAgent:
"""Build a minimal ``WorkflowAgent`` whose inner agent emits a fixed text."""
class _TextAgent(SupportsAgentRun):
def __init__(self, name: str, text: str) -> None:
self.id = str(uuid.uuid4())
self.name = name
self.description: str | None = None
self._text = text
def create_session(self, **kwargs: Any) -> AgentSession:
return AgentSession()
def get_session(self, *, service_session_id: str, **kwargs: Any) -> AgentSession:
return AgentSession()
@overload
def run(
self,
messages: Any = ...,
*,
stream: Literal[False] = ...,
session: AgentSession | None = ...,
**kwargs: Any,
) -> Awaitable[AgentResponse[Any]]: ...
@overload
def run(
self,
messages: Any = ...,
*,
stream: Literal[True],
session: AgentSession | None = ...,
**kwargs: Any,
) -> ResponseStream[AgentResponseUpdate, AgentResponse[Any]]: ...
def run(
self,
messages: Any = None,
*,
stream: bool = False,
session: AgentSession | None = None,
**kwargs: Any,
) -> Awaitable[AgentResponse] | ResponseStream[AgentResponseUpdate, AgentResponse]:
text = self._text
name = self.name
async def _aresult() -> AgentResponse:
return AgentResponse(messages=[Message("assistant", [Content.from_text(text=text)])])
async def _aiter() -> AsyncIterator[AgentResponseUpdate]:
yield AgentResponseUpdate(
contents=[Content.from_text(text=text)],
role="assistant",
author_name=name,
)
if stream:
return ResponseStream(_aiter(), finalizer=AgentResponse.from_updates)
return _aresult()
inner = _TextAgent("text-agent", text)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, inner).build()
return WorkflowAgent(workflow=workflow, name="Text Workflow Agent")
def _build_approval_workflow_agent(
*,
approval_request_id: str,
tool_name: str = "delete_file",
tool_arguments: dict[str, Any] | None = None,
final_text: str = "done",
) -> tuple[WorkflowAgent, _ToolApprovalWorkflowAgentMock]:
"""Build a ``WorkflowAgent`` whose inner agent emits a tool approval request."""
mock_agent = _ToolApprovalWorkflowAgentMock(
name="approval-agent",
tool_name=tool_name,
tool_arguments=tool_arguments or {"path": "/tmp/secret.txt"},
approval_request_ids=[approval_request_id],
final_text=final_text,
)
@executor
async def start(messages: list[Message], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
workflow = WorkflowBuilder(start_executor=start).add_edge(start, mock_agent).build()
workflow_agent = WorkflowAgent(workflow=workflow, name="Approval Workflow Agent")
return workflow_agent, mock_agent
class TestWorkflowAgentHosting:
"""End-to-end HTTP tests for ``ResponsesHostServer`` hosting a ``WorkflowAgent``.
These tests drive ``_handle_inner_workflow`` through the ASGI stack:
they exercise checkpoint write/restore (multi-turn) and the
tool-approval round-trip path, which is the primary differentiator
relative to the regular agent path.
"""
async def test_basic_text_response(self) -> None:
workflow_agent = _build_text_workflow_agent("hello from workflow")
server = _make_server(workflow_agent)
resp = await _post(server, input_text="hi", stream=False)
assert resp.status_code == 200
body = resp.json()
assert body["status"] == "completed"
text_found = any(
part.get("type") == "output_text" and part.get("text") == "hello from workflow"
for item in body["output"]
if item["type"] == "message"
for part in item.get("content", [])
)
assert text_found, f"Expected workflow output text in {body['output']}"
async def test_basic_text_response_streaming(self) -> None:
workflow_agent = _build_text_workflow_agent("hello stream")
server = _make_server(workflow_agent)
resp = await _post(server, input_text="hi", stream=True)
assert resp.status_code == 200
events = _parse_sse_events(resp.text)
types = _sse_event_types(events)
assert types[0] == "response.created"
assert types[-1] == "response.completed"
assert "response.output_text.delta" in types
text_done = [e for e in events if e["event"] == "response.output_text.done"]
assert any(e["data"]["text"] == "hello stream" for e in text_done)
async def test_non_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None:
workflow_agent, mock_agent = _build_approval_workflow_agent(approval_request_id="apr_wf_ns")
server = _make_server(workflow_agent)
resp = await _post(server, stream=False)
assert resp.status_code == 200
body = resp.json()
assert body["status"] == "completed"
approval_items = [it for it in body["output"] if it["type"] == "mcp_approval_request"]
assert len(approval_items) == 1
assert approval_items[0]["name"] == "delete_file"
assert approval_items[0]["server_label"] == "test_server"
approval_request_id = approval_items[0]["id"]
# The id surfaced over the wire is generated by the response stream
# builder; the original approval ``Content`` (carrying the inner
# ``function_call``) must be persisted under that id so the next
# turn can reconstruct it.
loaded = await server._approval_storage.load_approval_request( # pyright: ignore[reportPrivateUsage]
approval_request_id
)
assert loaded.type == "function_approval_request"
assert loaded.function_call.name == "delete_file" # type: ignore[attr-defined]
assert mock_agent.run_count == 1
async def test_streaming_emits_mcp_approval_request_and_persists_to_storage(self) -> None:
workflow_agent, mock_agent = _build_approval_workflow_agent(approval_request_id="apr_wf_st")
server = _make_server(workflow_agent)
resp = await _post(server, stream=True)
assert resp.status_code == 200
events = _parse_sse_events(resp.text)
types = _sse_event_types(events)
assert types[0] == "response.created"
assert types[-1] == "response.completed"
approval_request_id: str | None = None
for e in events:
if e["event"] != "response.output_item.added":
continue
item = e["data"].get("item") or {}
if item.get("type") == "mcp_approval_request":
approval_request_id = item.get("id")
break
assert approval_request_id is not None
loaded = await server._approval_storage.load_approval_request( # pyright: ignore[reportPrivateUsage]
approval_request_id
)
assert loaded.type == "function_approval_request"
assert mock_agent.run_count == 1
async def test_round_trip_approval_response_resumes_workflow_agent(self) -> None:
"""Two-turn HTTP round-trip:
Turn 1 emits ``mcp_approval_request`` and writes a workflow
checkpoint under the response id. Turn 2 sends the
``mcp_approval_response`` with ``previous_response_id`` set, so the
host restores the checkpoint, the WorkflowAgent routes the
approval response back to the paused inner agent, and the inner
agent emits the final assistant text.
"""
workflow_agent, mock_agent = _build_approval_workflow_agent(
approval_request_id="apr_wf_rt",
final_text="done with approval",
)
server = _make_server(workflow_agent)
first = await _post(server, stream=False)
assert first.status_code == 200
first_body = first.json()
first_response_id = first_body["id"]
approval_items = [it for it in first_body["output"] if it["type"] == "mcp_approval_request"]
assert len(approval_items) == 1
approval_request_id = approval_items[0]["id"]
assert mock_agent.run_count == 1
second_payload: dict[str, Any] = {
"model": "test-model",
"input": [
{
"type": "mcp_approval_response",
"approval_request_id": approval_request_id,
"approve": True,
}
],
"stream": False,
"previous_response_id": first_response_id,
}
second = await _post_json(server, second_payload)
assert second.status_code == 200
second_body = second.json()
assert second_body["status"] == "completed"
# The inner agent must have been resumed (restore replay + new turn).
# Restore call is a no-op for the mock (no input); the new-turn call
# delivers the approval response, so run_count grows by at least 1.
assert mock_agent.run_count >= 2
# The final assistant text from the resumed inner agent surfaces in
# the HTTP output.
text_pieces = [
part.get("text", "")
for item in second_body["output"]
if item["type"] == "message"
for part in item.get("content", [])
if part.get("type") == "output_text"
]
assert any("done with approval" in t for t in text_pieces), (
f"expected resumed workflow output, got {second_body['output']}"
)
# The new-turn invocation of the inner agent must have received the
# approval response routed back through WorkflowAgent.
approval_responses = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approval_responses) == 1
assert approval_responses[0].approved is True # type: ignore[attr-defined]
async def test_round_trip_approval_response_streaming(self) -> None:
"""Streaming variant of the round-trip: turn 2 is requested with
``stream=true`` and surfaces the resumed text as SSE events."""
workflow_agent, mock_agent = _build_approval_workflow_agent(
approval_request_id="apr_wf_rt_st",
final_text="streamed-done",
)
server = _make_server(workflow_agent)
first = await _post(server, stream=False)
first_body = first.json()
first_response_id = first_body["id"]
approval_request_id = next(it["id"] for it in first_body["output"] if it["type"] == "mcp_approval_request")
second = await _post_json(
server,
{
"model": "test-model",
"input": [
{
"type": "mcp_approval_response",
"approval_request_id": approval_request_id,
"approve": True,
}
],
"stream": True,
"previous_response_id": first_response_id,
},
)
assert second.status_code == 200
events = _parse_sse_events(second.text)
types = _sse_event_types(events)
assert types[0] == "response.created"
assert types[-1] == "response.completed"
text_done = [e for e in events if e["event"] == "response.output_text.done"]
assert any("streamed-done" in e["data"]["text"] for e in text_done)
assert mock_agent.run_count >= 2
async def test_round_trip_approval_response_rejected(self) -> None:
"""Sending ``approve=False`` must surface as ``approved=False`` to the
inner agent on resume."""
workflow_agent, mock_agent = _build_approval_workflow_agent(
approval_request_id="apr_wf_reject",
final_text="acknowledged",
)
server = _make_server(workflow_agent)
first = await _post(server, stream=False)
first_body = first.json()
first_response_id = first_body["id"]
approval_request_id = next(it["id"] for it in first_body["output"] if it["type"] == "mcp_approval_request")
second = await _post_json(
server,
{
"model": "test-model",
"input": [
{
"type": "mcp_approval_response",
"approval_request_id": approval_request_id,
"approve": False,
}
],
"stream": False,
"previous_response_id": first_response_id,
},
)
assert second.status_code == 200
approval_responses = [
c for m in mock_agent.last_run_messages for c in m.contents if c.type == "function_approval_response"
]
assert len(approval_responses) == 1
assert approval_responses[0].approved is False # type: ignore[attr-defined]
# endregion
+4
View File
@@ -10,6 +10,10 @@ pip install agent-framework-gemini --pre
The Gemini integration enables Microsoft Agent Framework applications to call Google Gemini models with familiar chat abstractions, including streaming, tool/function calling, and structured output.
## Structured Output
Gemini structured output can be configured with either a Pydantic model in `response_format`, a JSON schema mapping in `response_format`, or a Gemini-specific `response_schema`. Declarative agents that define `outputSchema` pass that schema through `response_format`.
## Authentication
The connector supports both `google-genai` authentication modes.
@@ -109,8 +109,8 @@ class GeminiChatOptions(ChatOptions[ResponseModelT], Generic[ResponseModelT], to
or ``types.Tool`` objects returned by ``get_code_interpreter_tool``, ``get_web_search_tool``,
``get_mcp_tool``, ``get_file_search_tool``, or ``get_maps_grounding_tool``.
tool_choice: How the model picks a tool. One of ``'auto'``, ``'none'``, or ``'required'``.
response_format: Pydantic model type for structured JSON output. The response text is
parsed into the model and exposed via ``ChatResponse.value``.
response_format: Pydantic model type or JSON schema mapping for structured JSON output.
The response text is parsed and exposed via ``ChatResponse.value``.
instructions: Extra system-level instructions prepended to the system message.
Not supported, and passing these raises a type error:
@@ -255,6 +255,29 @@ _OPTION_CONSUMED_KEYS: frozenset[str] = frozenset({
_OPTION_EXCLUDE_KEYS: frozenset[str] = _OPTION_EXPLICIT_KEYS | _OPTION_CONSUMED_KEYS
_JSON_SCHEMA_TYPES: frozenset[str] = frozenset({
"array",
"boolean",
"integer",
"null",
"number",
"object",
"string",
})
_JSON_SCHEMA_KEYWORDS: frozenset[str] = frozenset({
"$defs",
"additionalProperties",
"allOf",
"anyOf",
"enum",
"items",
"oneOf",
"properties",
"required",
"type",
})
_FINISH_REASON_MAP: dict[str, FinishReasonLiteral] = {
"STOP": "stop",
"MAX_TOKENS": "length",
@@ -747,9 +770,13 @@ class RawGeminiChatClient(
continue
kwargs[_OPTION_TRANSLATIONS.get(key, key)] = value
if options.get("response_format") or options.get("response_schema"):
response_format = options.get("response_format")
response_schema = options.get("response_schema")
if response_format is not None or response_schema is not None:
kwargs["response_mime_type"] = "application/json"
if schema := options.get("response_schema"):
if response_schema is not None:
kwargs["response_schema"] = response_schema
elif (schema := self._extract_response_schema(response_format)) is not None:
kwargs["response_schema"] = schema
if tools := self._prepare_tools(options):
kwargs["tools"] = tools
@@ -762,6 +789,48 @@ class RawGeminiChatClient(
return types.GenerateContentConfig(**kwargs)
@staticmethod
def _extract_response_schema(response_format: Any) -> dict[str, Any] | None:
"""Extract a Gemini response schema from supported mapping response_format shapes."""
if not isinstance(response_format, Mapping):
return None
mapping = cast("Mapping[str, Any]", response_format)
if (nested := RawGeminiChatClient._extract_response_schema(mapping.get("format"))) is not None:
return nested
json_schema = mapping.get("json_schema")
if isinstance(json_schema, Mapping):
schema = cast("Mapping[str, Any]", json_schema).get("schema")
if isinstance(schema, Mapping):
return dict(cast("Mapping[str, Any]", schema))
schema = mapping.get("schema")
if isinstance(schema, Mapping):
return dict(cast("Mapping[str, Any]", schema))
if RawGeminiChatClient._is_json_schema_mapping(mapping):
return dict(mapping)
return None
@staticmethod
def _is_json_schema_mapping(value: Mapping[str, Any]) -> bool:
"""Return True when a mapping appears to be a JSON Schema rather than a response-format envelope."""
if not any(keyword in value for keyword in _JSON_SCHEMA_KEYWORDS):
return False
schema_type = value.get("type")
if schema_type is None:
return True
if isinstance(schema_type, str):
return schema_type in _JSON_SCHEMA_TYPES
if isinstance(schema_type, Sequence) and not isinstance(schema_type, (str, bytes)):
entries = cast("Sequence[object]", schema_type)
return all(isinstance(item, str) and item in _JSON_SCHEMA_TYPES for item in entries)
return False
def _prepare_tools(self, options: Mapping[str, Any]) -> list[types.Tool] | None:
"""Translate the framework tool list into Gemini API tool objects.
@@ -9,7 +9,7 @@ from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import Content, FunctionTool, Message
from agent_framework import Agent, Content, FunctionTool, Message
from google.genai import types
from pydantic import BaseModel
@@ -915,6 +915,20 @@ async def test_response_format_populates_value_on_chat_response() -> None:
assert response.value == Reply(text="hello")
async def test_response_format_mapping_populates_value_on_chat_response() -> None:
"""When response_format is a JSON schema mapping, ChatResponse.value must parse the response text."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"text": "hello"}')]))
schema = {"type": "object", "properties": {"text": {"type": "string"}}}
response = await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": schema},
)
assert response.value == {"text": "hello"}
async def test_response_schema_added_to_config() -> None:
"""Sets both response_mime_type and the raw schema on the config when response_schema is given."""
client, mock = _make_gemini_client()
@@ -931,6 +945,284 @@ async def test_response_schema_added_to_config() -> None:
assert config.response_schema == schema
async def test_response_format_raw_json_schema_added_to_config() -> None:
"""For declarative outputSchema, response_format may already be a raw JSON schema mapping."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
schema = {
"type": "object",
"properties": {"answer": {"type": "string", "description": "The answer."}},
"required": ["answer"],
}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": schema},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema == schema
async def test_agent_default_options_response_format_raw_schema_added_to_config() -> None:
"""Agent default_options is the path used by declarative outputSchema."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]}
agent = Agent(client=client, default_options={"response_format": schema})
await agent.run("Hi")
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema == schema
async def test_response_format_complex_raw_json_schema_preserved() -> None:
"""Nested declarative schemas should be forwarded without losing shape or constraints."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "ok"}')]))
schema = {
"type": "object",
"properties": {
"answer": {"type": "string"},
"citations": {
"type": "array",
"items": {
"type": "object",
"properties": {
"source": {"type": "string"},
"confidence": {"type": "number"},
},
"required": ["source"],
"additionalProperties": False,
},
},
},
"required": ["answer"],
"additionalProperties": False,
}
await client.get_response(
messages=[
Message(
role="user",
contents=[
Content.from_text(
"Summarize a long document while preserving citation metadata.\n" + ("context\n" * 128)
)
],
)
],
options={"response_format": schema},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_schema == schema
async def test_response_format_json_schema_envelope_added_to_config() -> None:
"""OpenAI-style json_schema envelopes should still provide Gemini with the inner schema."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": {"type": "json_schema", "json_schema": {"name": "Answer", "schema": schema}}},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema == schema
async def test_response_format_format_envelope_added_to_config() -> None:
"""Responses-style format envelopes should also provide Gemini with the nested schema."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": {"format": {"type": "json_schema", "name": "Answer", "schema": schema}}},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema == schema
async def test_response_format_direct_schema_key_added_to_config() -> None:
"""Provider-normalized mappings with a direct schema key should be accepted."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": {"schema": schema}},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema == schema
async def test_response_format_json_schema_envelope_preserves_empty_schema() -> None:
"""An explicitly empty JSON schema is still a schema and should not be dropped as falsy."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
schema: dict[str, Any] = {}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": {"type": "json_schema", "json_schema": {"name": "AnyJson", "schema": schema}}},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_schema == schema
async def test_response_format_anyof_raw_schema_added_to_config() -> None:
"""Raw schemas without a type should still be recognized when they use JSON Schema keywords."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='"ok"')]))
schema = {"anyOf": [{"type": "string"}, {"type": "number"}]}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": schema},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_schema == schema
async def test_response_format_union_type_raw_schema_added_to_config() -> None:
"""JSON Schema union type arrays should be treated as raw schemas."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "hello"}')]))
schema = {"type": ["object", "null"], "properties": {"answer": {"type": "string"}}}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": schema},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_schema == schema
async def test_response_format_json_object_does_not_set_schema() -> None:
"""A JSON-object response_format requests JSON output but is not itself a Gemini response schema."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": {"type": "json_object"}},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema is None
async def test_response_format_json_schema_without_inner_schema_does_not_set_schema() -> None:
"""A json_schema envelope without a schema should not be mistaken for a raw JSON schema."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": {"type": "json_schema", "json_schema": {"name": "MissingSchema"}}},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema is None
async def test_response_schema_takes_precedence_over_response_format_schema() -> None:
"""An explicit Gemini response_schema should win when both schema options are present."""
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text="{}")]))
response_format_schema = {"type": "object", "properties": {"name": {"type": "string"}}}
response_schema = {"type": "object", "properties": {"id": {"type": "integer"}}}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": response_format_schema, "response_schema": response_schema},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_schema == response_schema
async def test_response_format_raw_schema_kept_with_tools() -> None:
"""Structured output must still reach Gemini when function tools are present."""
def calculator(expression: str) -> str:
"""Evaluate a simple expression."""
return expression
tool = FunctionTool(name="calculator", func=calculator)
client, mock = _make_gemini_client()
mock.aio.models.generate_content = AsyncMock(return_value=_make_response([_make_part(text='{"answer": "4"}')]))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}, "required": ["answer"]}
await client.get_response(
messages=[Message(role="user", contents=[Content.from_text("What is 2 + 2?")])],
options={"tools": [tool], "response_format": schema},
)
config: types.GenerateContentConfig = mock.aio.models.generate_content.call_args.kwargs["config"]
assert config.response_schema == schema
assert config.tools is not None
assert config.tools[0].function_declarations[0].name == "calculator"
async def test_streaming_response_format_raw_schema_added_to_config() -> None:
"""Streaming requests use the same config path and should also forward raw schema mappings."""
client, mock = _make_gemini_client()
chunks = [_make_response([_make_part(text='{"answer": "hello"}')], finish_reason="STOP")]
mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
stream = client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": schema},
stream=True,
)
async for _ in stream:
pass
config: types.GenerateContentConfig = mock.aio.models.generate_content_stream.call_args.kwargs["config"]
assert config.response_mime_type == "application/json"
assert config.response_schema == schema
async def test_streaming_response_format_mapping_populates_final_value() -> None:
"""Streaming responses should preserve mapping response_format for final value parsing."""
client, mock = _make_gemini_client()
chunks = [_make_response([_make_part(text='{"answer": "hello"}')], finish_reason="STOP")]
mock.aio.models.generate_content_stream = AsyncMock(return_value=_async_iter(chunks))
schema = {"type": "object", "properties": {"answer": {"type": "string"}}}
stream = client.get_response(
messages=[Message(role="user", contents=[Content.from_text("Hi")])],
options={"response_format": schema},
stream=True,
)
async for _ in stream:
pass
final = await stream.get_final_response()
assert final.value == {"answer": "hello"}
async def test_streaming_response_format_passed_to_build_response_stream() -> None:
"""Verifies that response_format is forwarded to _build_response_stream when streaming
so that structured output parsing works correctly on the final assembled response.
@@ -8,29 +8,34 @@ This module provides ``Mem0ContextProvider``, built on the new
from __future__ import annotations
import asyncio
import logging
import sys
from collections.abc import Awaitable
from contextlib import AbstractAsyncContextManager
from typing import TYPE_CHECKING, Any, ClassVar
from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypedDict
from agent_framework import Message
from agent_framework._sessions import AgentSession, ContextProvider, SessionContext
from mem0 import AsyncMemory, AsyncMemoryClient
if sys.version_info >= (3, 11):
from typing import NotRequired, Self, TypedDict # pragma: no cover
from typing import Self # pragma: no cover
else:
from typing_extensions import NotRequired, Self, TypedDict # pragma: no cover
from typing_extensions import Self # pragma: no cover
if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
class _MemorySearchResponse_v1_1(TypedDict):
results: list[dict[str, Any]]
relations: NotRequired[list[dict[str, Any]]]
logger = logging.getLogger(__name__)
MemoryRecord: TypeAlias = dict[str, object]
_MemorySearchResponse_v2 = list[dict[str, Any]]
class SearchResults(TypedDict):
results: list[MemoryRecord]
SearchResponse: TypeAlias = list[MemoryRecord] | SearchResults
class Mem0ContextProvider(ContextProvider):
@@ -106,28 +111,85 @@ class Mem0ContextProvider(ContextProvider):
if not input_text.strip():
return
filters = self._build_filters()
# Query entity partitions independently to bypass strict logical AND limitations
# Mem0 OSS and Platform SDKs expose inconsistent search typings.
search_tasks: list[Awaitable[Any]] = []
# AsyncMemory (OSS) expects user_id/agent_id/run_id as direct kwargs
# AsyncMemoryClient (Platform) expects them in a filters dict
search_kwargs: dict[str, Any] = {"query": input_text}
if isinstance(self.mem0_client, AsyncMemory):
search_kwargs.update(filters)
else:
search_kwargs["filters"] = filters
# 1. Query User partition independently
if self.user_id:
user_kwargs = self._build_search_kwargs(input_text, "user_id", self.user_id)
search_tasks.append(self.mem0_client.search(**user_kwargs)) # type: ignore[reportUnknownMemberType, reportUnknownArgumentType]
search_response: _MemorySearchResponse_v1_1 | _MemorySearchResponse_v2 = await self.mem0_client.search( # type: ignore[misc]
**search_kwargs,
)
# 2. Query Agent partition independently
if self.agent_id:
agent_kwargs = self._build_search_kwargs(input_text, "agent_id", self.agent_id)
search_tasks.append(self.mem0_client.search(**agent_kwargs)) # type: ignore[reportUnknownMemberType, reportUnknownArgumentType]
if isinstance(search_response, list):
memories = search_response
elif isinstance(search_response, dict) and "results" in search_response:
memories = search_response["results"]
else:
memories = [search_response]
# Fall back to an app-scoped search when only application_id is configured
if not search_tasks and self.application_id:
app_kwargs: dict[str, Any] = {"query": input_text}
if isinstance(self.mem0_client, AsyncMemory):
app_kwargs["app_id"] = self.application_id
else:
app_kwargs["filters"] = {"app_id": self.application_id}
search_tasks.append(self.mem0_client.search(**app_kwargs)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
if not search_tasks:
return
line_separated_memories = "\n".join(memory.get("memory", "") for memory in memories)
results: list[SearchResponse | BaseException] = await asyncio.gather(*search_tasks, return_exceptions=True)
# Merge and deduplicate results
memories: list[MemoryRecord] = []
seen_memory_ids: set[str] = set()
failed_tasks_count: int = 0
for search_response in results:
if isinstance(search_response, asyncio.CancelledError):
raise search_response
if isinstance(search_response, BaseException):
failed_tasks_count += 1
logger.error(
"Mem0 partition search task failed: %s",
search_response,
exc_info=(type(search_response), search_response, search_response.__traceback__),
)
continue
current_memories: list[MemoryRecord] = []
if isinstance(search_response, list):
current_memories = [mem for mem in search_response if isinstance(mem, dict)]
elif isinstance(search_response, dict):
results_field = search_response.get("results")
if isinstance(results_field, list):
current_memories = [
item
for item in results_field
if isinstance(item, dict) # pyright: ignore[reportUnknownVariableType]
]
else:
logger.warning(
"Unexpected Mem0 search response format: %s",
type(results_field).__name__,
)
for mem in current_memories:
mem_id = mem.get("id")
if mem_id is not None and not isinstance(mem_id, str):
mem_id = str(mem_id)
if mem_id is not None and mem_id in seen_memory_ids:
continue
if mem_id is not None:
seen_memory_ids.add(mem_id)
memories.append(mem)
if failed_tasks_count == len(search_tasks):
logger.error("All Mem0 retrieval tasks failed. Context provider is unable to verify memory state.")
line_separated_memories = "\n".join(str(memory.get("memory", "")) for memory in memories)
if line_separated_memories:
context.extend_messages(
self.source_id,
@@ -159,12 +221,21 @@ class Mem0ContextProvider(ContextProvider):
]
if messages:
await self.mem0_client.add( # type: ignore[misc]
messages=messages,
user_id=self.user_id,
agent_id=self.agent_id,
metadata={"application_id": self.application_id},
)
add_kwargs: dict[str, Any] = {
"messages": messages,
"user_id": self.user_id,
"agent_id": self.agent_id,
}
# Inject the application scope using the matching signature format for each SDK variant
if isinstance(self.mem0_client, AsyncMemory):
if self.application_id:
add_kwargs["app_id"] = self.application_id
else:
if self.application_id:
add_kwargs["filters"] = {"app_id": self.application_id}
await self.mem0_client.add(**add_kwargs) # type: ignore[misc, call-arg]
# -- Internal methods ------------------------------------------------------
@@ -173,15 +244,21 @@ class Mem0ContextProvider(ContextProvider):
if not self.agent_id and not self.user_id and not self.application_id:
raise ValueError("At least one of the filters: agent_id, user_id, or application_id is required.")
def _build_filters(self) -> dict[str, Any]:
"""Build search filters from initialization parameters."""
filters: dict[str, Any] = {}
if self.user_id:
filters["user_id"] = self.user_id
if self.agent_id:
filters["agent_id"] = self.agent_id
if self.application_id:
filters["app_id"] = self.application_id
def _build_search_kwargs(self, input_text: str, entity_key: str, entity_value: str) -> dict[str, Any]:
"""Build search keyword arguments formatted for OSS vs Platform clients."""
filters: dict[str, Any] = {"query": input_text}
if isinstance(self.mem0_client, AsyncMemory):
# AsyncMemory (OSS) expects direct kwargs
filters[entity_key] = entity_value
if self.application_id:
filters["app_id"] = self.application_id
else:
# AsyncMemoryClient (Platform) expects a filters dict
filters["filters"] = {entity_key: entity_value}
if self.application_id:
filters["filters"]["app_id"] = self.application_id
return filters
@@ -3,7 +3,7 @@
from __future__ import annotations
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import AgentResponse, Message
@@ -193,39 +193,59 @@ class TestBeforeRun:
assert call_kwargs["user_id"] == "u1"
assert "filters" not in call_kwargs
async def test_oss_client_all_scoping_params(self, mock_oss_mem0_client: AsyncMock) -> None:
"""OSS client with all scoping parameters passes them as direct kwargs."""
@pytest.mark.asyncio
async def test_oss_client_all_scoping_params_except_app_id(self, mock_oss_mem0_client: AsyncMock) -> None:
"""OSS client with all scoping parameters passes them as isolated concurrent kwargs."""
mock_oss_mem0_client.search.return_value = []
provider = Mem0ContextProvider(
source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1", agent_id="a1", application_id="app1"
source_id="mem0",
mem0_client=mock_oss_mem0_client,
user_id="u1",
agent_id="a1"
)
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=["Hello"])], session_id="s1")
mock_context = MagicMock(spec=SessionContext)
mock_msg = MagicMock()
mock_msg.text = "hello"
mock_context.input_messages = [mock_msg]
mock_context.response = None
await provider.before_run(
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
) # type: ignore[arg-type]
agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={}
)
call_kwargs = mock_oss_mem0_client.search.call_args.kwargs
assert call_kwargs["user_id"] == "u1"
assert call_kwargs["agent_id"] == "a1"
assert "filters" not in call_kwargs
# Re-aligned assertion: We expect 2 separate concurrent calls instead of 1 combined call
assert mock_oss_mem0_client.search.call_count == 2
mock_oss_mem0_client.search.assert_any_call(query="hello", user_id="u1")
mock_oss_mem0_client.search.assert_any_call(query="hello", agent_id="a1")
async def test_platform_client_passes_filters_dict(self, mock_mem0_client: AsyncMock) -> None:
"""Platform AsyncMemoryClient should receive scoping params in a filters dict."""
@pytest.mark.asyncio
async def test_platform_client_passes_filters_dict_except_app_id(self, mock_mem0_client: AsyncMock) -> None:
"""Platform client passes scoping parameters concurrently inside the nested filters dictionary."""
mock_mem0_client.search.return_value = []
provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client, user_id="u1")
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=["Hello"])], session_id="s1")
provider = Mem0ContextProvider(
source_id="mem0",
mem0_client=mock_mem0_client,
user_id="u1",
agent_id="a1",
)
mock_context = MagicMock(spec=SessionContext)
mock_msg = MagicMock()
mock_msg.text = "hello"
mock_context.input_messages = [mock_msg]
mock_context.response = None
await provider.before_run(
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
) # type: ignore[arg-type]
agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={}
)
call_kwargs = mock_mem0_client.search.call_args.kwargs
assert call_kwargs["query"] == "Hello"
assert "filters" in call_kwargs
assert call_kwargs["filters"]["user_id"] == "u1"
# Re-aligned assertion: Platform client isolates filters per call to bypass AND limitations
assert mock_mem0_client.search.call_count == 2
mock_mem0_client.search.assert_any_call(query="hello", filters={"user_id": "u1"})
mock_mem0_client.search.assert_any_call(query="hello", filters={"agent_id": "a1"})
# -- after_run tests -----------------------------------------------------------
@@ -318,8 +338,8 @@ class TestAfterRun:
with pytest.raises(ValueError, match="At least one of the filters"):
await provider.after_run(agent=None, session=session, context=ctx, state=session.state) # type: ignore[arg-type]
async def test_stores_with_application_id_metadata(self, mock_mem0_client: AsyncMock) -> None:
"""application_id is passed in metadata."""
async def test_stores_with_application_id_filters(self, mock_mem0_client: AsyncMock) -> None:
"""application_id is passed in filters."""
provider = Mem0ContextProvider(
source_id="mem0", mem0_client=mock_mem0_client, user_id="u1", application_id="app1"
)
@@ -331,7 +351,7 @@ class TestAfterRun:
agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
) # type: ignore[arg-type]
assert mock_mem0_client.add.call_args.kwargs["metadata"] == {"application_id": "app1"}
assert mock_mem0_client.add.call_args.kwargs["filters"] == {"app_id": "app1"}
# -- _validate_filters tests --------------------------------------------------
@@ -358,15 +378,20 @@ class TestValidateFilters:
provider._validate_filters()
# -- _build_filters tests -----------------------------------------------------
# -- _build_search_kwargs tests -----------------------------------------------------
class TestBuildFilters:
"""Test _build_filters method."""
class TestBuildSearchKwargs:
"""Test _build_search_kwargs method."""
def test_user_id_only(self, mock_mem0_client: AsyncMock) -> None:
provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client, user_id="u1")
assert provider._build_filters() == {"user_id": "u1"}
# Pass the 3 required arguments
result = provider._build_search_kwargs("test query", "user_id", "u1")
# AsyncMock triggers the Platform client nested 'filters' structure
assert result == {"query": "test query", "filters": {"user_id": "u1"}}
def test_all_params(self, mock_mem0_client: AsyncMock) -> None:
provider = Mem0ContextProvider(
@@ -376,28 +401,66 @@ class TestBuildFilters:
agent_id="a1",
application_id="app1",
)
assert provider._build_filters() == {
"user_id": "u1",
"agent_id": "a1",
"app_id": "app1",
# Test that app_id correctly merges with the isolated target entity
result = provider._build_search_kwargs("test query", "agent_id", "a1")
assert result == {
"query": "test query",
"filters": {
"agent_id": "a1",
"app_id": "app1",
},
}
def test_excludes_none_values(self, mock_mem0_client: AsyncMock) -> None:
provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client, user_id="u1")
filters = provider._build_filters()
assert "agent_id" not in filters
assert "run_id" not in filters
assert "app_id" not in filters
# application_id is None by default, it should not appear in the dictionary
result = provider._build_search_kwargs("test query", "user_id", "u1")
assert "app_id" not in result.get("filters", {})
def test_no_run_id_in_search_filters(self, mock_mem0_client: AsyncMock) -> None:
"""run_id is excluded from search filters so memories work across sessions."""
provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client, user_id="u1")
filters = provider._build_filters()
assert "run_id" not in filters
result = provider._build_search_kwargs("test query", "user_id", "u1")
assert "run_id" not in result.get("filters", {})
assert "run_id" not in result
def test_empty_when_no_params(self, mock_mem0_client: AsyncMock) -> None:
# Validates base query payload generation
provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client)
assert provider._build_filters() == {}
result = provider._build_search_kwargs("test query", "custom_key", "custom_val")
assert result == {"query": "test query", "filters": {"custom_key": "custom_val"}}
@pytest.mark.asyncio
async def test_before_run_application_only_fallback(self, mock_mem0_client: AsyncMock) -> None:
provider = Mem0ContextProvider(
source_id="mem0", mem0_client=mock_mem0_client, application_id="app_fallback_test"
)
# Mock a valid message list and session container setup
mock_context = MagicMock(spec=SessionContext)
mock_msg = MagicMock()
mock_msg.text = "Retrieve systemic fallback memory traces"
mock_context.input_messages = [mock_msg]
mock_context.response = None
mock_mem0_client.search = AsyncMock(return_value=[{"id": "m1", "memory": "System configuration template"}])
await provider.before_run(
agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={}
)
# Verify that an application-scoped search task executed successfully
assert mock_mem0_client.search.call_count == 1
mock_context.extend_messages.assert_called_once()
# -- Context manager tests -----------------------------------------------------
+1
View File
@@ -7,6 +7,7 @@
"**/demos/**",
"**/_to_delete/**",
"**/05-end-to-end/**",
"**/harness/**",
"**/agent_with_foundry_tracing.py",
"**/azure_responses_client_with_foundry.py"
],
+1
View File
@@ -7,6 +7,7 @@
"**/demos/**",
"**/_to_delete/**",
"**/05-end-to-end/**",
"**/harness/**",
"**/agent_with_foundry_tracing.py",
"**/azure_responses_client_with_foundry.py",
"**/github_copilot/**"
@@ -0,0 +1,94 @@
# Harness Console
A Textual-based terminal UI for running and observing AI agents built with the Agent Framework.
## Quick Start
```python
from console import run_agent_async, build_default_observers
await run_agent_async(
agent=my_agent,
session=my_session,
observers=build_default_observers(),
)
```
See [`harness_research.py`](../harness_research.py) for a complete example.
## Package Structure
```
console/
├── __init__.py # Public API exports
├── harness_console.py # run_agent_async() entry point
├── app.py # HarnessApp (Textual application)
├── app_state.py # HarnessAppState, enums, data types
├── agent_runner.py # HarnessAgentRunner (streaming orchestration)
├── state_driver.py # IUXStateDriver protocol
├── textual_state_driver.py # Textual implementation of IUXStateDriver
├── formatters.py # Tool call formatters
├── observers/ # Lifecycle observers
│ ├── base.py # ConsoleObserver abstract base
│ ├── text_output.py # Streaming text display
│ ├── tool_call_display.py # Tool call formatting
│ ├── tool_approval.py # User approval for tool calls
│ ├── error_display.py # Error messages
│ ├── usage_display.py # Token usage tracking
│ └── reasoning_display.py # Reasoning/thinking blocks
├── components/ # Textual UI widgets
│ ├── scroll_panel.py # Conversation history
│ ├── text_input.py # User text input
│ ├── list_selection.py # Multiple choice selector
│ ├── agent_status.py # Spinner + usage display
│ └── agent_mode_help.py # Mode indicator + help text
└── commands/ # Slash command handlers
├── base.py # CommandHandler abstract base
├── exit_handler.py # /exit
├── mode_handler.py # /mode [plan|execute]
├── todo_handler.py # /todos
└── session_handler.py # /session-export, /session-import
```
## Public API
| Export | Description |
|--------|-------------|
| `run_agent_async` | Main entry point — runs the Textual app with an agent |
| `build_default_observers` | Factory for the standard observer set |
| `build_default_command_handlers` | Factory for slash command handlers |
| `ConsoleObserver` | Base class for custom observers |
| `ToolCallFormatter` | Base class for custom tool formatters |
| `CommandHandler` | Base class for custom slash commands |
## Architecture
The console follows a unidirectional data flow:
```
AgentRunner → Observers → StateDriver → AppState → Textual UI
User Input (app.py)
```
- **AgentRunner** streams responses from the agent and dispatches events to observers.
- **Observers** process events (text chunks, tool calls, errors) and update the state driver.
- **StateDriver** (`IUXStateDriver`) mutates `HarnessAppState` and notifies the UI.
- **Textual App** reads state and syncs widgets on each notification.
### Key Design Choices
| Concern | Approach |
|---------|----------|
| Rendering | Textual widgets + Rich markup (no manual ANSI) |
| State | Single `HarnessAppState` dataclass, mutated by driver |
| Streaming text | Truncate-and-rewrite on RichLog for flicker-free updates |
| Extensibility | Custom observers, formatters, and commands via base classes |
| Follow-up questions | Observer returns `FollowUpQuestion` → UI shows prompt/choices |
## Dependencies
- `textual` — TUI framework
- `rich` — Text formatting
- `agent-framework` — Core agent framework
@@ -0,0 +1,27 @@
# Copyright (c) Microsoft. All rights reserved.
"""Harness Console - A Textual-based TUI for AI agent interactions.
This package provides a rich terminal interface for running and observing
AI agents, with streaming output, tool call display, follow-up questions,
and token usage tracking.
"""
from .commands import CommandHandler, build_default_command_handlers
from .formatters import ToolCallFormatter
from .harness_console import run_agent_async
from .observers import (
ConsoleObserver,
build_default_observers,
build_observers_with_planning,
)
__all__ = [
"CommandHandler",
"ConsoleObserver",
"ToolCallFormatter",
"build_default_command_handlers",
"build_default_observers",
"build_observers_with_planning",
"run_agent_async",
]
@@ -0,0 +1,343 @@
# Copyright (c) Microsoft. All rights reserved.
"""Agent runner orchestration for the harness console.
This module provides the HarnessAgentRunner class, which orchestrates agent
invocations with observer lifecycle management. It handles:
- User input dispatch
- Agent streaming with observer notifications
- Follow-up action collection
- Streaming state management
"""
from __future__ import annotations
import asyncio
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from agent_framework import Agent, AgentSession
from .app_state import FollowUpAction
from .observers.base import ConsoleObserver
from .state_driver import IUXStateDriver
class HarnessAgentRunner:
"""Orchestrates agent invocations driven by user-input events from the UI.
The component invokes the runner's input handlers (run_turn) directly;
the runner mutates UI state through the supplied IUXStateDriver.
This is a minimal implementation focusing on the core agent loop without
command handling or complex message injection (those can be added later).
"""
def __init__(
self,
agent: Agent,
observers: list[ConsoleObserver],
state_driver: IUXStateDriver,
*,
max_context_window_tokens: int | None = None,
max_output_tokens: int | None = None,
) -> None:
"""Initialize the agent runner.
Args:
agent: The agent to orchestrate.
observers: List of console observers for lifecycle events.
state_driver: The UI state driver for observer updates.
max_context_window_tokens: Optional max context window size for usage display.
max_output_tokens: Optional max output tokens for usage display.
"""
self._agent = agent
self._observers = observers
self._ux = state_driver
self._max_context_window_tokens = max_context_window_tokens
self._max_output_tokens = max_output_tokens
self._input_gate = asyncio.Semaphore(1) # Single turn at a time
async def run_turn(
self,
user_input: str,
session: AgentSession | None = None,
) -> None:
"""Run a single agent turn with the given user input.
Echoes the input, then delegates to the agent loop.
Args:
user_input: The user's input text.
session: Optional agent session for conversation history.
"""
async with self._input_gate:
self._ux.write_user_input_echo(user_input)
from agent_framework import Message
messages = [Message(role="user", contents=[user_input])]
await self._run_agent_loop(messages, session)
async def start_agent_turn(
self,
messages: list,
session: AgentSession | None = None,
) -> None:
"""Resume the agent loop with pre-built messages (from follow-up responses).
Called by the app after the user finishes answering follow-up questions.
If messages is empty, just completes the turn.
Args:
messages: List of Message objects to send to the agent.
session: Optional agent session.
"""
async with self._input_gate:
if not messages:
self._complete_turn()
return
await self._run_agent_loop(messages, session)
async def _run_agent_loop(
self,
messages: list,
session: AgentSession | None,
) -> None:
"""Run the agent loop, re-invoking as needed for follow-up messages.
Loops while there are messages to send. After each stream:
- Collects follow-up actions from observers
- If questions exist queue them and return (UI will collect answers)
- If only direct messages loop with those messages
- If nothing complete the turn
Args:
messages: Initial messages to send.
session: Optional agent session.
"""
next_messages = messages
while next_messages:
# Configure run options
options = self._configure_run_options(session)
# Begin streaming
self._ux.begin_streaming()
self._ux.begin_streaming_output()
self._ux.set_show_spinner(True)
try:
await self._stream_response_messages(next_messages, session, options)
except Exception as ex:
self._ux.append_info_line(
f"❌ Stream error: {ex.__class__.__name__}:\n{ex}",
color="red",
)
# Stop spinner and end streaming output
self._ux.set_show_spinner(False)
# Collect follow-up actions from observers
follow_up_actions = await self._collect_follow_up_actions(session)
# Separate direct messages from questions
has_follow_ups = len(follow_up_actions) > 0
# Write no-text warning if applicable
await self._ux.write_no_text_warning(has_follow_ups)
# Enqueue all follow-up actions
for action in follow_up_actions:
self._ux.enqueue_follow_up_action(action)
# Check if there are pending questions (UI needs user input)
if self._ux.has_pending_questions():
# Pause — the UI will collect answers and call start_agent_turn
return
# No questions — drain any accumulated direct messages and loop
drained = self._ux.take_follow_up_responses()
next_messages = drained if drained else None
self._complete_turn()
def _complete_turn(self) -> None:
"""Complete the current turn (end streaming)."""
self._ux.end_streaming()
def _configure_run_options(
self,
session: AgentSession | None,
) -> dict:
"""Configure run options via observers.
Each observer can modify the options dict to influence agent behavior.
Args:
session: Optional agent session.
Returns:
Options dict for agent.run().
"""
options = {}
for observer in self._observers:
observer.configure_run_options(options, self._agent, session)
return options
async def _stream_response(
self,
user_input: str,
session: AgentSession | None,
options: dict,
) -> None:
"""Stream agent response from a text input and dispatch to observers.
Args:
user_input: The user's input text.
session: Optional agent session.
options: Run options configured by observers.
"""
# Stream response using agent.run(stream=True)
stream = self._agent.run(
user_input,
stream=True,
session=session,
options=options,
)
# Process each update chunk
async for update in stream:
await self._dispatch_update(update, session)
# Extract usage from the final response
self._extract_usage(stream)
async def _stream_response_messages(
self,
messages: list,
session: AgentSession | None,
options: dict,
) -> None:
"""Stream agent response from Message objects and dispatch to observers.
Args:
messages: List of Message objects to send.
session: Optional agent session.
options: Run options configured by observers.
"""
stream = self._agent.run(
messages,
stream=True,
session=session,
options=options,
)
async for update in stream:
await self._dispatch_update(update, session)
self._extract_usage(stream)
def _extract_usage(self, stream) -> None:
"""Extract token usage from a completed stream."""
try:
get_final = getattr(stream, "get_final_response", None)
if not get_final:
return
import inspect
if inspect.iscoroutinefunction(get_final):
return
final_response = get_final()
if final_response is None:
return
usage = getattr(final_response, "usage_details", None)
if not isinstance(usage, dict):
return
input_tokens = usage.get("input_token_count", 0) or 0
output_tokens = usage.get("output_token_count", 0) or 0
if input_tokens or output_tokens:
self._ux.set_usage_text(self._format_usage(input_tokens, output_tokens))
except (AttributeError, TypeError):
pass
async def _dispatch_update(
self,
update, # AgentResponseUpdate
session: AgentSession | None,
) -> None:
"""Dispatch a single update to all observers.
Calls observer lifecycle methods in order:
1. on_response_update (once per update)
2. on_content (for each content item)
3. on_text (if text is present)
Args:
update: The agent response update.
session: Optional agent session.
"""
# on_response_update
for observer in self._observers:
await observer.on_response_update(self._ux, update, self._agent, session)
# on_content for each content item
if hasattr(update, "contents") and update.contents:
for content in update.contents:
for observer in self._observers:
await observer.on_content(self._ux, content, self._agent, session)
# on_text for text chunks
if hasattr(update, "text") and update.text:
for observer in self._observers:
await observer.on_text(self._ux, update.text, self._agent, session)
async def _collect_follow_up_actions(
self,
session: AgentSession | None,
) -> list[FollowUpAction]:
"""Collect follow-up actions from all observers.
Called after streaming completes to gather any follow-up questions
or messages from observers.
Args:
session: Optional agent session.
Returns:
List of follow-up actions from all observers.
"""
actions: list[FollowUpAction] = []
for observer in self._observers:
observer_actions = await observer.on_stream_complete(
self._ux, self._agent, session
)
if observer_actions:
actions.extend(observer_actions)
return actions
def _format_usage(self, input_tokens: int, output_tokens: int) -> str:
"""Format token counts matching C# harness style: 📊 Tokens — input: X | output: Y | total: Z."""
total_tokens = input_tokens + output_tokens
input_budget = None
if self._max_context_window_tokens and self._max_output_tokens:
input_budget = self._max_context_window_tokens - self._max_output_tokens
return (
f"📊 Tokens — input: {self._format_token_count(input_tokens, input_budget)}"
f" | output: {self._format_token_count(output_tokens, self._max_output_tokens)}"
f" | total: {self._format_token_count(total_tokens, self._max_context_window_tokens)}"
)
@staticmethod
def _format_token_count(count: int, budget: int | None) -> str:
"""Format a token count, optionally showing budget percentage."""
if budget and budget > 0:
pct = count / budget * 100
return f"{count:,}/{budget:,} ({pct:.1f}%)"
return f"{count:,}"
@@ -0,0 +1,541 @@
# Copyright (c) Microsoft. All rights reserved.
"""Main Textual application for the harness console.
This module provides the HarnessApp - the main Textual application that
composes all UI components and integrates with the agent runner.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from textual import on, work
from textual.app import App, ComposeResult
from textual.binding import Binding
from textual.containers import Container, Vertical
from textual.css.query import NoMatches
from textual.widgets import Input, Static
from .app_state import (
BottomPanelMode,
HarnessAppState,
OutputEntryType,
)
from .components import (
AgentModeAndHelp,
AgentStatus,
HarnessListSelection,
HarnessScrollPanel,
HarnessTextInput,
PromptRule,
)
from .textual_state_driver import HarnessConsoleUXStateDriver
if TYPE_CHECKING:
from agent_framework import Agent, AgentSession
from .agent_runner import HarnessAgentRunner
from .commands import CommandHandler
from .observers.base import ConsoleObserver
class HarnessApp(App[None]):
"""Main Textual application for the harness console.
Composes the scroll panel (conversation history), status bar (spinner, usage),
mode/help display, and bottom panel (text input, list selection, or streaming
indicator). Routes user input to the agent runner.
"""
CSS = """
Screen {
background: $background;
}
#scroll-panel {
height: 1fr;
padding: 0 1;
background: transparent;
}
#bottom-panel {
height: auto;
}
#text-input-container {
height: 1;
display: block;
}
#list-selection-container {
height: auto;
max-height: 12;
display: none;
}
#streaming-indicator {
height: 1;
display: none;
}
#status-bar {
height: 1;
}
#mode-help {
height: 1;
}
#top-rule {
height: 1;
}
#bottom-rule {
height: 1;
}
#separator-rule {
height: 1;
}
#text-input {
height: 1;
}
.hidden {
display: none;
}
.visible {
display: block;
}
.input-field {
border: none;
padding: 0;
min-height: 1;
height: 1;
background: transparent;
}
.input-field:focus {
border: none;
background: transparent;
}
.prompt-container {
height: 1;
}
.prompt-label {
width: 2;
min-width: 2;
height: 1;
}
"""
BINDINGS = [
Binding("ctrl+c", "quit", "Quit", show=False),
Binding("ctrl+q", "quit", "Quit", show=False),
]
def __init__(
self,
agent: Agent,
observers: list[ConsoleObserver],
session: AgentSession | None = None,
mode_colors: dict[str, str] | None = None,
initial_mode: str | None = None,
placeholder: str = "Type a message and press Enter...",
title: str = "Harness Console",
max_context_window_tokens: int | None = None,
max_output_tokens: int | None = None,
command_handlers: list[CommandHandler] | None = None,
) -> None:
"""Initialize the harness console application.
Args:
agent: The agent to run.
observers: List of console observers.
session: Optional agent session.
mode_colors: Optional mode color mapping.
initial_mode: Initial agent mode.
placeholder: Input placeholder text.
title: Application title.
max_context_window_tokens: Optional max context window tokens for usage display.
max_output_tokens: Optional max output tokens for usage display.
command_handlers: Optional list of command handlers. If None, auto-detected.
"""
super().__init__()
self.title = title
self._agent = agent
self._observers = observers
self._session = session
self._mode_colors = mode_colors
self._initial_mode = initial_mode
self._placeholder = placeholder
self._max_context_window_tokens = max_context_window_tokens
self._max_output_tokens = max_output_tokens
# Build command handlers
if command_handlers is None:
from .commands import build_default_command_handlers
self._command_handlers = build_default_command_handlers(
agent, mode_colors=mode_colors
)
else:
self._command_handlers = command_handlers
# Compute help text from command handlers
help_parts = [
h.get_help_text()
for h in self._command_handlers
if h.get_help_text() is not None
]
help_text = ", ".join(help_parts) if help_parts else None
# State and driver
self._app_state = HarnessAppState(
placeholder=placeholder,
mode_text=initial_mode,
help_text=help_text,
)
self._ux_driver = HarnessConsoleUXStateDriver(
app_state=self._app_state,
on_state_changed=self._on_state_changed,
mode_colors=mode_colors,
)
# Agent runner (created after init)
self._runner: HarnessAgentRunner | None = None
@property
def ux_driver(self) -> HarnessConsoleUXStateDriver:
"""Get the UX state driver."""
return self._ux_driver
@property
def runner(self) -> HarnessAgentRunner | None:
"""Get the agent runner."""
return self._runner
def compose(self) -> ComposeResult:
"""Compose the application layout."""
with Vertical():
# Main scroll panel for conversation history
yield HarnessScrollPanel(id="scroll-panel")
# Blank line separating scroll content from status area
yield Static(" ", id="separator-rule")
# Status bar (spinner + usage)
yield AgentStatus(id="status-bar")
# Top rule (mode-colored)
yield PromptRule(id="top-rule")
# Bottom panel - switches between text input, list selection, streaming
with Container(id="bottom-panel"):
# Text input (default)
with Container(id="text-input-container"):
text_input = HarnessTextInput(id="text-input")
text_input.placeholder = self._placeholder
yield text_input
# List selection (for follow-up questions)
with Container(id="list-selection-container"):
yield HarnessListSelection(id="list-selection")
# Bottom rule (mode-colored)
yield PromptRule(id="bottom-rule")
# Mode and help
yield AgentModeAndHelp(id="mode-help")
def on_mount(self) -> None:
"""Initialize after mount."""
# Create agent runner now that everything is set up
from .agent_runner import HarnessAgentRunner
self._runner = HarnessAgentRunner(
agent=self._agent,
observers=self._observers,
state_driver=self._ux_driver,
max_context_window_tokens=self._max_context_window_tokens,
max_output_tokens=self._max_output_tokens,
)
# Set initial mode
if self._initial_mode:
self._ux_driver.current_mode = self._initial_mode
# Focus the text input
try:
text_input = self.query_one("#text-input", HarnessTextInput)
text_input.focus_input()
except NoMatches:
pass
# Set initial rule colors and mode display
self._sync_mode_help()
# --- Event handlers ---
@on(HarnessTextInput.Submitted)
def on_text_submitted(self, event: HarnessTextInput.Submitted) -> None:
"""Handle text input submission."""
text = event.value.strip()
if not text:
return
if self._app_state.pending_questions:
# Answer the current follow-up question
self._handle_follow_up_answer(text)
elif self._app_state.mode == BottomPanelMode.STREAMING:
# Input during streaming (message injection placeholder)
pass
elif text.startswith("/"):
# Try command handlers
self._try_command_handlers(text)
else:
# Normal user input - run agent turn
self._run_agent_turn(text)
@work(exclusive=True, thread=False)
async def _try_command_handlers(self, text: str) -> None:
"""Try each command handler; fall through to agent if none match."""
session = self._session
if session is None:
# No session — fall through to agent turn
self._run_agent_turn(text)
return
for handler in self._command_handlers:
if await handler.try_handle(text, session, self._ux_driver):
# Command handled — check for shutdown/session swap signals
self._process_command_signals()
return
# No handler matched — treat as normal agent input
self._run_agent_turn(text)
def _process_command_signals(self) -> None:
"""Check and process signals set by command handlers."""
if self._app_state.shutdown_requested:
self.exit()
return
if self._app_state.replaced_session is not None:
self._session = self._app_state.replaced_session # type: ignore[assignment]
self._app_state.replaced_session = None
self._ux_driver.append_info_line("Session replaced.")
self._sync_ui_from_state()
@on(HarnessListSelection.Selected)
def on_list_selected(self, event: HarnessListSelection.Selected) -> None:
"""Handle list selection."""
self._handle_follow_up_answer(event.value)
# --- Agent turn ---
@work(exclusive=True, thread=False)
async def _run_agent_turn(self, text: str) -> None:
"""Run an agent turn in a background worker."""
if self._runner is None:
return
await self._runner.run_turn(text, session=self._session)
# After turn completes, check for follow-up questions
self._sync_ui_from_state()
# --- Follow-up question handling ---
@work(exclusive=True, thread=False)
async def _handle_follow_up_answer(self, answer: str) -> None:
"""Handle a user's answer to a follow-up question."""
if not self._app_state.pending_questions:
return
question = self._app_state.pending_questions[0]
# Call the continuation
result_message = await question.continuation(answer, self._ux_driver)
# Add result to accumulated responses
if result_message is not None:
self._ux_driver.add_follow_up_response(result_message)
# Advance to next question
self._ux_driver.advance_follow_up_question()
# If no more questions, resume the agent with accumulated responses
if not self._app_state.pending_questions:
responses = self._ux_driver.take_follow_up_responses()
if responses and self._runner:
await self._runner.start_agent_turn(responses, session=self._session)
self._sync_ui_from_state()
# --- State synchronization ---
def _on_state_changed(self) -> None:
"""Called by state driver when state changes - schedule UI sync.
Since the agent runner uses @work(thread=False), state changes happen
on the main event loop. We use call_later to batch updates.
"""
self.call_later(self._sync_ui_from_state)
def _sync_ui_from_state(self) -> None:
"""Synchronize UI components with current application state."""
state = self._app_state
# Update scroll panel with new entries
self._sync_scroll_panel()
# Update bottom panel mode
self._sync_bottom_panel(state.mode)
# Hide status bar and mode/help during list selection (matching C#)
is_list_mode = state.mode == BottomPanelMode.LIST_SELECTION
self._sync_chrome_visibility(not is_list_mode)
# Update status bar
self._sync_status_bar()
# Update mode/help display
self._sync_mode_help()
def _sync_scroll_panel(self) -> None:
"""Sync the scroll panel with output entries."""
try:
panel = self.query_one("#scroll-panel", HarnessScrollPanel)
except NoMatches:
return
entries = self._app_state.output_entries
rendered_count = getattr(self, "_rendered_entry_count", 0)
if rendered_count < len(entries):
# There are new entries to render
for entry in entries[rendered_count:]:
if entry.type == OutputEntryType.STREAMING_TEXT:
panel.set_streaming_entry(entry)
else:
# End any active streaming before appending other entry types
panel.end_streaming()
panel.append_entry(entry)
self._rendered_entry_count = len(entries)
elif rendered_count == len(entries) and entries:
# Same count — check if the last entry is a streaming entry that was mutated
last_entry = entries[-1]
if last_entry.type == OutputEntryType.STREAMING_TEXT:
panel.set_streaming_entry(last_entry)
def _sync_bottom_panel(self, mode: BottomPanelMode) -> None:
"""Switch the bottom panel between text input, list, and streaming."""
try:
text_container = self.query_one("#text-input-container")
list_container = self.query_one("#list-selection-container")
except NoMatches:
return
if mode == BottomPanelMode.TEXT_INPUT:
text_container.display = True
list_container.display = False
# Restore focus to text input
try:
text_input = self.query_one("#text-input", HarnessTextInput)
text_input.focus_input()
except NoMatches:
pass
elif mode == BottomPanelMode.LIST_SELECTION:
text_container.display = False
list_container.display = True
self._sync_list_selection()
elif mode == BottomPanelMode.STREAMING:
text_container.display = True
list_container.display = False
def _sync_list_selection(self) -> None:
"""Sync the list selection widget with state."""
try:
list_widget = self.query_one("#list-selection", HarnessListSelection)
except NoMatches:
return
state = self._app_state
list_widget.title = state.list_selection_title or ""
list_widget.options = list(state.list_selection_options)
list_widget.allow_custom_text = state.list_selection_custom_text_placeholder is not None
if state.list_selection_custom_text_placeholder:
try:
custom_input = list_widget.query_one("#custom-input", Input)
custom_input.placeholder = state.list_selection_custom_text_placeholder
except Exception:
pass
# Focus the option list so keyboard navigation works immediately
list_widget.focus_list()
def _sync_status_bar(self) -> None:
"""Sync the status bar with state."""
try:
status = self.query_one("#status-bar", AgentStatus)
except NoMatches:
return
state = self._app_state
status.show_spinner = state.show_spinner
status.usage_text = state.usage_text or ""
def _sync_mode_help(self) -> None:
"""Sync the mode/help display and rule colors with state."""
try:
mode_help = self.query_one("#mode-help", AgentModeAndHelp)
except NoMatches:
return
state = self._app_state
mode_help.mode = state.mode_text or ""
mode_help.mode_color = state.mode_color or "blue"
mode_help.help_text = state.help_text or ""
# Sync rule colors to match mode
color = state.mode_color or "cyan"
try:
top_rule = self.query_one("#top-rule", PromptRule)
top_rule.rule_color = color
except NoMatches:
pass
try:
bottom_rule = self.query_one("#bottom-rule", PromptRule)
bottom_rule.rule_color = color
except NoMatches:
pass
def _sync_chrome_visibility(self, visible: bool) -> None:
"""Show or hide chrome elements (status bar, mode/help).
During list selection mode, these are hidden to give more vertical
space to the scroll panel and list picker.
Args:
visible: Whether chrome elements should be visible.
"""
import contextlib
with contextlib.suppress(NoMatches):
self.query_one("#status-bar", AgentStatus).display = visible
with contextlib.suppress(NoMatches):
self.query_one("#mode-help", AgentModeAndHelp).display = visible
# --- Rendering count tracking ---
_rendered_entry_count: int = 0
@@ -0,0 +1,260 @@
# Copyright (c) Microsoft. All rights reserved.
"""Application state and core data types for the harness console.
This module defines enums, dataclasses, follow-up action types, and the
HarnessAppState dataclass which holds all UI state that may change during
application execution. The state driver mutates this state to coordinate
between the agent runner and the Textual UI components.
"""
from __future__ import annotations
from collections.abc import Awaitable, Callable
from dataclasses import dataclass, field
from enum import Enum
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from agent_framework import Message
from .state_driver import IUXStateDriver
# region Enums
class OutputEntryType(Enum):
"""Type of output entry in the console conversation."""
USER_INPUT = "user_input"
"""User input echo (e.g., 'You: hello')."""
STREAMING_TEXT = "streaming_text"
"""In-progress streaming text from the agent (accumulated chunk by chunk)."""
INFO_LINE = "info_line"
"""Informational line (tool calls, errors, usage, approval requests, etc.)."""
STREAM_FOOTER = "stream_footer"
"""Stream footer (e.g., '(no text response from agent)')."""
PENDING_MESSAGE = "pending_message"
"""Pending injected message notification."""
class BottomPanelMode(Enum):
"""Mode of the bottom panel UI."""
TEXT_INPUT = "text_input"
"""Show text input for user messages."""
LIST_SELECTION = "list_selection"
"""Show choice list for user selection."""
STREAMING = "streaming"
"""Show 'streaming...' indicator while agent is generating."""
# endregion
# region Output Entry
@dataclass
class OutputEntry:
"""A single output entry in the console conversation history.
Used internally by the state driver to track conversation output,
including streaming text, tool calls, errors, and user input echoes.
Args:
type: The type of output entry.
text: The text content of the entry.
color: Optional Rich color string (e.g., "cyan", "red", "dim").
"""
type: OutputEntryType
text: str
color: str | None = None
# endregion
# region Follow-Up Actions
class FollowUpAction:
"""Base class for follow-up actions returned by observers.
Follow-up actions describe either a question to ask the user
(via FollowUpQuestion subclasses) or a message to add directly
to the next agent input (FollowUpMessage).
"""
pass
@dataclass
class FollowUpQuestion(FollowUpAction):
"""A question to ask the user with a continuation.
The continuation delegate is invoked with the user's answer and the
UX state driver, and returns an optional Message to add to the next
agent invocation.
Args:
prompt: The question text shown to the user.
continuation: Async function invoked with the user's answer and state driver.
Returns an optional Message to add to the next agent input.
"""
prompt: str
continuation: Callable[[str, IUXStateDriver], Awaitable[Message | None]]
@dataclass
class TextFollowUpQuestion(FollowUpQuestion):
"""A free-form text question.
The user may type any response. This is the base FollowUpQuestion type
with no additional constraints.
"""
pass
@dataclass
class ChoiceFollowUpQuestion(FollowUpQuestion):
"""A multiple choice question.
The user picks from the provided choices, with an optional ability to
enter custom text when allow_custom_text is True.
Args:
prompt: The question text shown to the user.
choices: List of pre-defined choices.
allow_custom_text: If True, the user may type a custom response in
addition to the listed choices.
continuation: Async function invoked with the user's choice/text and
state driver. Returns an optional Message to add to the next agent input.
"""
choices: list[str]
allow_custom_text: bool = False
@dataclass
class FollowUpMessage(FollowUpAction):
"""A message to add directly to the next agent invocation without prompting.
Used when an observer wants to inject a message into the conversation
without user interaction (e.g., automatic tool results, system messages).
Args:
message: The Message to add to the conversation.
"""
message: Message
# endregion
# region Application State
@dataclass
class HarnessAppState:
"""All UI state for the harness console application.
This state is mutated by the UX state driver and read by the Textual
app to update the UI.
"""
# --- Bottom panel mode ---
mode: BottomPanelMode = BottomPanelMode.TEXT_INPUT
"""Which component is shown in the bottom panel."""
# --- Follow-up question queue ---
pending_questions: list[FollowUpQuestion] = field(default_factory=list)
"""Queue of follow-up questions waiting for user answers.
The head ([0]) is the question currently being displayed; subsequent items
are dispatched in order as each is answered.
"""
accumulated_follow_up_responses: list[Message] = field(default_factory=list)
"""Accumulated follow-up response messages collected during the current agent turn.
Both direct FollowUpMessages emitted by observers and continuation results
from answered questions. Consumed by the runner via take_follow_up_responses().
"""
# --- Text input (active in TextInput / Streaming modes) ---
prompt: str = "> "
"""The prompt string for text input mode."""
placeholder: str = ""
"""Placeholder text shown when the input is empty."""
input_text: str = ""
"""The current input text being typed."""
input_enabled: bool = True
"""Whether input is enabled (disabled during streaming without injection)."""
streaming_prompt: str = "(agent is running...)"
"""The prompt to show during streaming when input is disabled."""
# --- List selection (active in ListSelection mode) ---
list_selection_title: str | None = None
"""Title text displayed above the list selection."""
list_selection_options: list[str] = field(default_factory=list)
"""The list selection options."""
list_selection_index: int = 0
"""The highlighted option index in list selection mode."""
list_selection_custom_text_placeholder: str | None = None
"""Placeholder text for the custom text input option in the list."""
list_selection_custom_input_text: str = ""
"""Current text being typed into the list's custom text option."""
# --- Scroll / output area ---
output_entries: list[OutputEntry] = field(default_factory=list)
"""Output entries in the scroll area conversation history."""
queued_items: list[str] = field(default_factory=list)
"""Queued input items to display (pending injected messages)."""
# --- Agent mode + status display ---
mode_color: str | None = None
"""Rich color string for the rule borders and mode label."""
mode_text: str | None = None
"""Current mode name displayed (e.g., 'plan', 'execute')."""
help_text: str | None = None
"""Help text displayed below the bottom rule (available commands)."""
show_spinner: bool = False
"""Whether the agent status spinner is visible."""
usage_text: str | None = None
"""Formatted token usage text to display in the status bar."""
# --- Command handler signals ---
shutdown_requested: bool = False
"""Set to True when /exit is invoked; the app should exit."""
replaced_session: object | None = None
"""When set, the app should swap its session to this AgentSession."""
@@ -0,0 +1,65 @@
# Copyright (c) Microsoft. All rights reserved.
"""Command handler package for the harness console.
Provides slash-command handling (e.g., /exit, /mode, /todos, /session-export)
that intercepts user input before it reaches the agent.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from .base import CommandHandler
from .exit_handler import ExitCommandHandler
from .mode_handler import ModeCommandHandler
from .session_handler import SessionCommandHandler
from .todo_handler import TodoCommandHandler
if TYPE_CHECKING:
from agent_framework import Agent
__all__ = [
"CommandHandler",
"ExitCommandHandler",
"ModeCommandHandler",
"SessionCommandHandler",
"TodoCommandHandler",
"build_default_command_handlers",
]
def build_default_command_handlers(
agent: Agent,
*,
mode_colors: dict[str, str] | None = None,
) -> list[CommandHandler]:
"""Build the default set of command handlers by inspecting the agent.
Auto-detects TodoProvider and AgentModeProvider from the agent's
context_providers list.
Args:
agent: The agent to inspect for providers.
mode_colors: Optional mapping of mode names to Rich color strings.
Returns:
List of command handlers in evaluation order.
"""
from agent_framework import AgentModeProvider, TodoProvider
todo_provider: TodoProvider | None = None
mode_provider: AgentModeProvider | None = None
for provider in getattr(agent, "context_providers", []):
if isinstance(provider, TodoProvider) and todo_provider is None:
todo_provider = provider
elif isinstance(provider, AgentModeProvider) and mode_provider is None:
mode_provider = provider
return [
ExitCommandHandler(),
TodoCommandHandler(todo_provider),
ModeCommandHandler(mode_provider, mode_colors),
SessionCommandHandler(),
]
@@ -0,0 +1,58 @@
# Copyright (c) Microsoft. All rights reserved.
"""Abstract base class for console command handlers.
Command handlers intercept user input starting with '/' and execute
local commands before input reaches the agent. They are checked in order;
the first handler that accepts the input prevents further handlers from
being checked.
"""
from __future__ import annotations
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from agent_framework import AgentSession
from ..state_driver import IUXStateDriver
class CommandHandler(ABC):
"""Base class for console command handlers.
Subclasses implement get_help_text() for the mode bar and
try_handle() to intercept matching commands.
"""
@abstractmethod
def get_help_text(self) -> str | None:
"""Get the help text for this command.
Displayed in the mode-and-help bar. Return None if the
command is not currently available.
Returns:
Help text like '/todos (show todo list)', or None.
"""
...
@abstractmethod
async def try_handle(
self,
user_input: str,
session: AgentSession,
ux: IUXStateDriver,
) -> bool:
"""Attempt to handle the given user input.
Args:
user_input: The raw user input string.
session: The current agent session.
ux: The UX state driver for rendering output.
Returns:
True if this handler handled the input; False otherwise.
"""
...
@@ -0,0 +1,35 @@
# Copyright (c) Microsoft. All rights reserved.
"""Exit command handler — /exit to quit the console."""
from __future__ import annotations
from typing import TYPE_CHECKING
from .base import CommandHandler
if TYPE_CHECKING:
from agent_framework import AgentSession
from ..state_driver import IUXStateDriver
class ExitCommandHandler(CommandHandler):
"""Handle the /exit command to shut down the console application."""
def get_help_text(self) -> str | None:
"""Return help text for the exit command."""
return "/exit (quit)"
async def try_handle(
self,
user_input: str,
session: AgentSession,
ux: IUXStateDriver,
) -> bool:
"""Handle /exit by requesting shutdown."""
if user_input.strip().lower() != "/exit":
return False
ux.request_shutdown()
return True
@@ -0,0 +1,81 @@
# Copyright (c) Microsoft. All rights reserved.
"""Mode command handler — /mode to show or switch agent mode."""
from __future__ import annotations
from typing import TYPE_CHECKING
from .base import CommandHandler
if TYPE_CHECKING:
from agent_framework import AgentModeProvider, AgentSession
from ..state_driver import IUXStateDriver
class ModeCommandHandler(CommandHandler):
"""Handle the /mode command to display or switch the current agent mode."""
def __init__(
self,
mode_provider: AgentModeProvider | None,
mode_colors: dict[str, str] | None = None,
) -> None:
"""Initialize with mode provider and color mapping.
Args:
mode_provider: The mode provider, or None if not available.
mode_colors: Optional mapping of mode names to Rich color strings.
"""
self._mode_provider = mode_provider
self._mode_colors = mode_colors or {}
def get_help_text(self) -> str | None:
"""Return help text, or None if mode provider is unavailable."""
if self._mode_provider is None:
return None
return "/mode [plan|execute] (show or switch mode)"
async def try_handle(
self,
user_input: str,
session: AgentSession,
ux: IUXStateDriver,
) -> bool:
"""Handle /mode [name] command."""
stripped = user_input.strip()
lower = stripped.lower()
if not (lower == "/mode" or lower.startswith("/mode ")):
return False
if self._mode_provider is None:
ux.append_info_line("AgentModeProvider is not available.")
return True
parts = stripped.split(None, 1)
if len(parts) < 2:
# Show current mode
from agent_framework import get_agent_mode
current = get_agent_mode(session)
ux.append_info_line(f"Current mode: {current}")
return True
# Switch mode
new_mode = parts[1].strip()
try:
from agent_framework import set_agent_mode
normalized = set_agent_mode(session, new_mode)
color = self._mode_colors.get(normalized)
ux.set_mode(normalized, color)
ux.append_info_line(
f"Switched to {normalized} mode.",
color=color,
)
except ValueError as ex:
ux.append_info_line(str(ex), color="red")
return True
@@ -0,0 +1,107 @@
# Copyright (c) Microsoft. All rights reserved.
"""Session command handler — /session-export and /session-import."""
from __future__ import annotations
import json
from typing import TYPE_CHECKING
from .base import CommandHandler
if TYPE_CHECKING:
from agent_framework import AgentSession
from ..state_driver import IUXStateDriver
class SessionCommandHandler(CommandHandler):
"""Handle /session-export and /session-import commands."""
def get_help_text(self) -> str | None:
"""Return help text for session commands."""
return "/session-export <file> | /session-import <file>"
async def try_handle(
self,
user_input: str,
session: AgentSession,
ux: IUXStateDriver,
) -> bool:
"""Handle session export/import commands."""
stripped = user_input.strip()
command = stripped.split(None, 1)[0].lower() if stripped else ""
if command == "/session-export":
await self._handle_export(stripped, session, ux)
return True
if command == "/session-import":
await self._handle_import(stripped, ux)
return True
return False
async def _handle_export(
self,
user_input: str,
session: AgentSession,
ux: IUXStateDriver,
) -> None:
"""Export the current session to a JSON file."""
parts = user_input.split(None, 1)
if len(parts) < 2:
ux.append_info_line("Usage: /session-export <filename>")
return
filename = parts[1].strip()
try:
serialized = session.to_dict()
json_str = json.dumps(serialized, indent=2)
self._write_file(filename, json_str)
ux.append_info_line(f"Session exported to {filename}")
except Exception as ex:
ux.append_info_line(
f"Failed to export session to {filename}: {ex}",
color="red",
)
async def _handle_import(
self,
user_input: str,
ux: IUXStateDriver,
) -> None:
"""Import a session from a JSON file."""
parts = user_input.split(None, 1)
if len(parts) < 2:
ux.append_info_line("Usage: /session-import <filename>")
return
filename = parts[1].strip()
try:
from agent_framework import AgentSession
json_str = self._read_file(filename)
data = json.loads(json_str)
new_session = AgentSession.from_dict(data)
ux.replace_session(new_session)
ux.append_info_line(f"Session imported from {filename}")
except FileNotFoundError:
ux.append_info_line(f"File not found: {filename}", color="red")
except Exception as ex:
ux.append_info_line(
f"Failed to import session from {filename}: {ex}",
color="red",
)
@staticmethod
def _write_file(filename: str, content: str) -> None:
"""Write content to a file (sync helper to satisfy ASYNC230)."""
with open(filename, "w", encoding="utf-8") as f: # noqa: ASYNC230
f.write(content)
@staticmethod
def _read_file(filename: str) -> str:
"""Read content from a file (sync helper to satisfy ASYNC230)."""
with open(filename, encoding="utf-8") as f: # noqa: ASYNC230
return f.read()
@@ -0,0 +1,66 @@
# Copyright (c) Microsoft. All rights reserved.
"""Todo command handler — /todos to display the todo list."""
from __future__ import annotations
from typing import TYPE_CHECKING
from .base import CommandHandler
if TYPE_CHECKING:
from agent_framework import AgentSession, TodoProvider
from ..state_driver import IUXStateDriver
class TodoCommandHandler(CommandHandler):
"""Handle the /todos command to display the current todo list."""
def __init__(self, todo_provider: TodoProvider | None) -> None:
"""Initialize with the todo provider.
Args:
todo_provider: The todo provider, or None if not available.
"""
self._todo_provider = todo_provider
def get_help_text(self) -> str | None:
"""Return help text, or None if todo provider is unavailable."""
if self._todo_provider is None:
return None
return "/todos (show todo list)"
async def try_handle(
self,
user_input: str,
session: AgentSession,
ux: IUXStateDriver,
) -> bool:
"""Handle /todos by displaying the todo list."""
if user_input.strip().lower() != "/todos":
return False
if self._todo_provider is None:
ux.append_info_line("TodoProvider is not available.")
return True
todos = await self._todo_provider.store.load_items(
session, source_id=self._todo_provider.source_id
)
if not todos:
ux.append_info_line("No todos yet.")
return True
ux.append_info_line("── Todo List ──")
for item in todos:
status = "" if item.is_complete else ""
color = "dim" if item.is_complete else None
description = f"{item.description}" if item.description else ""
ux.append_info_line(
f"[{status}] #{item.id} {item.title}{description}",
color=color,
)
return True
@@ -0,0 +1,23 @@
# Copyright (c) Microsoft. All rights reserved.
"""UI components for the harness console.
This module provides Textual widgets for building the harness console UI,
including status displays, input fields, choice selectors, and scrolling panels.
"""
from .agent_status import AgentStatus
from .list_selection import HarnessListSelection
from .mode_help import AgentModeAndHelp
from .prompt_rule import PromptRule
from .scroll_panel import HarnessScrollPanel
from .text_input import HarnessTextInput
__all__ = [
"AgentStatus",
"AgentModeAndHelp",
"HarnessListSelection",
"PromptRule",
"HarnessScrollPanel",
"HarnessTextInput",
]
@@ -0,0 +1,66 @@
# Copyright (c) Microsoft. All rights reserved.
"""Agent status widget with spinner animation and usage statistics."""
from __future__ import annotations
from textual.reactive import reactive
from textual.widgets import Static
class AgentStatus(Static):
"""Agent status bar with animated spinner and token usage display.
Displays an animated braille pattern spinner when the agent is active,
along with token usage statistics. The component automatically updates
the spinner animation at ~10fps for smooth visual feedback.
Attributes:
show_spinner: Whether to display the animated spinner.
usage_text: Token usage text to display (e.g., "1.2K in / 856 out").
"""
# Braille pattern spinner frames for smooth animation
SPINNER_FRAMES = ["", "", "", "", "", "", "", "", "", ""]
show_spinner: reactive[bool] = reactive(False)
usage_text: reactive[str] = reactive("")
def __init__(self, **kwargs) -> None:
"""Initialize the agent status widget."""
super().__init__(**kwargs)
self._spinner_index = 0
def on_mount(self) -> None:
"""Start the spinner animation timer when the widget is mounted."""
# Update spinner at ~10fps (every 0.1 seconds)
self.set_interval(0.1, self._advance_spinner)
def _advance_spinner(self) -> None:
"""Advance the spinner to the next frame."""
if self.show_spinner:
self._spinner_index = (self._spinner_index + 1) % len(self.SPINNER_FRAMES)
self.refresh()
def render(self) -> str:
"""Render the status bar with spinner and usage text.
Returns:
Formatted string with Rich markup for spinner and usage display.
"""
if not self.show_spinner and not self.usage_text:
return ""
parts = []
if self.show_spinner:
frame = self.SPINNER_FRAMES[self._spinner_index]
parts.append(f"[cyan]{frame}[/cyan]")
else:
# Keep consistent spacing when spinner is off
parts.append(" ")
if self.usage_text:
parts.append(f"[dim]{self.usage_text}[/dim]")
return " ".join(parts)
@@ -0,0 +1,269 @@
# Copyright (c) Microsoft. All rights reserved.
"""List selection widget with optional custom text input."""
from __future__ import annotations
from textual import on
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Container
from textual.css.query import NoMatches
from textual.events import Key
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
from textual.widgets import Input, Label, OptionList
from textual.widgets.option_list import Option
class HarnessListSelection(Widget):
"""List selection widget with numbered choices and optional custom text input.
Displays a title, a list of numbered choices that can be selected via
keyboard navigation or number keys (1-9), and an optional custom text
input field at the bottom.
All child nodes (title label, option list, custom input) are always
present in the DOM; visibility is toggled via reactive watchers.
Navigation:
- Down arrow on last list item moves focus to the custom text input
- Up arrow on the custom text input moves focus back to the option list
- When custom input has focus, the option list highlight is cleared
Attributes:
title: The title text displayed above the options.
options: List of option strings to display.
allow_custom_text: Whether to show a custom text input field.
"""
DEFAULT_CSS = """
HarnessListSelection {
height: auto;
max-height: 12;
}
HarnessListSelection .list-selection-container {
height: auto;
}
HarnessListSelection #selection-title {
height: auto;
color: $text;
text-style: bold;
padding: 0 0 0 0;
}
HarnessListSelection #option-list {
height: auto;
max-height: 8;
border: none;
padding: 0;
}
HarnessListSelection #custom-input {
height: auto;
min-height: 1;
margin-top: 0;
border: tall transparent;
}
HarnessListSelection #custom-input:focus {
border: tall $accent;
}
"""
BINDINGS = [
Binding("1", "select_option(0)", "Select option 1", show=False),
Binding("2", "select_option(1)", "Select option 2", show=False),
Binding("3", "select_option(2)", "Select option 3", show=False),
Binding("4", "select_option(3)", "Select option 4", show=False),
Binding("5", "select_option(4)", "Select option 5", show=False),
Binding("6", "select_option(5)", "Select option 6", show=False),
Binding("7", "select_option(6)", "Select option 7", show=False),
Binding("8", "select_option(7)", "Select option 8", show=False),
Binding("9", "select_option(8)", "Select option 9", show=False),
]
title: reactive[str] = reactive("")
options: reactive[list[str]] = reactive(list, always_update=True)
allow_custom_text: reactive[bool] = reactive(False)
class Selected(Message):
"""Message sent when an option is selected.
Attributes:
value: The selected option text or custom text.
"""
def __init__(self, value: str) -> None:
"""Initialize the Selected message.
Args:
value: The selected option text or custom text.
"""
self.value = value
super().__init__()
def compose(self) -> ComposeResult:
"""Compose the widget — all nodes are always present.
Yields:
Title label (hidden if empty), option list, custom input (hidden by default).
"""
with Container(classes="list-selection-container"):
yield Label("", id="selection-title")
yield OptionList(id="option-list")
yield Input(
placeholder="Or type a custom response...",
id="custom-input",
)
def on_mount(self) -> None:
"""Configure initial visibility after mount."""
title_label = self.query_one("#selection-title", Label)
title_label.display = bool(self.title)
custom_input = self.query_one("#custom-input", Input)
custom_input.display = self.allow_custom_text
self._update_options()
def on_key(self, event: Key) -> None:
"""Handle key navigation between option list and custom input.
Args:
event: The key event.
"""
if not self.allow_custom_text:
return
option_list = self.query_one("#option-list", OptionList)
custom_input = self.query_one("#custom-input", Input)
# Down arrow on last item → move to custom input
if event.key == "down" and option_list.has_focus:
last_index = option_list.option_count - 1
if last_index >= 0 and option_list.highlighted == last_index:
option_list.highlighted = None # type: ignore[assignment]
custom_input.focus()
event.prevent_default()
event.stop()
# Up arrow on custom input → move back to option list (last item)
elif event.key == "up" and custom_input.has_focus:
last_index = option_list.option_count - 1
if last_index >= 0:
option_list.highlighted = last_index
option_list.focus()
event.prevent_default()
event.stop()
@on(Input.Changed, "#custom-input")
def on_custom_input_focused_or_changed(self, event: Input.Changed) -> None:
"""Clear option list highlight when user is typing in custom input.
Args:
event: The input changed event.
"""
option_list = self.query_one("#option-list", OptionList)
option_list.highlighted = None # type: ignore[assignment]
def watch_title(self, new_title: str) -> None:
"""Update the title label when the title changes.
Args:
new_title: The new title text.
"""
try:
label = self.query_one("#selection-title", Label)
label.update(new_title)
label.display = bool(new_title)
except NoMatches:
pass
def watch_options(self, new_options: list[str]) -> None:
"""Update the option list when options change.
Args:
new_options: The new list of options.
"""
import contextlib
with contextlib.suppress(NoMatches):
self._update_options()
def watch_allow_custom_text(self, allow: bool) -> None:
"""Show/hide the custom input field.
Args:
allow: Whether to show the custom text input.
"""
try:
custom_input = self.query_one("#custom-input", Input)
custom_input.display = allow
except NoMatches:
pass
def _update_options(self) -> None:
"""Update the OptionList with numbered options."""
try:
option_list = self.query_one("#option-list", OptionList)
option_list.clear_options()
for i, option_text in enumerate(self.options):
display_text = f"{i + 1}. {option_text}" if i < 9 else f" {option_text}"
option_list.add_option(Option(display_text, id=str(i)))
except NoMatches:
pass
@on(OptionList.OptionSelected)
def on_option_selected(self, event: OptionList.OptionSelected) -> None:
"""Handle option selection from the list.
Args:
event: The OptionList.OptionSelected event.
"""
option_index = int(event.option.id or "0")
if 0 <= option_index < len(self.options):
selected_value = self.options[option_index]
self.post_message(self.Selected(selected_value))
@on(Input.Submitted)
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle custom text input submission.
Args:
event: The Input.Submitted event.
"""
if self.allow_custom_text and event.value:
self.post_message(self.Selected(event.value))
event.input.clear()
def action_select_option(self, index: int) -> None:
"""Select an option by index (0-based).
Args:
index: The option index to select.
"""
if 0 <= index < len(self.options):
selected_value = self.options[index]
self.post_message(self.Selected(selected_value))
def focus_list(self) -> None:
"""Focus the option list."""
try:
option_list = self.query_one("#option-list", OptionList)
option_list.focus()
except NoMatches:
pass
def focus_custom_input(self) -> None:
"""Focus the custom text input field."""
if self.allow_custom_text:
try:
custom_input = self.query_one("#custom-input", Input)
custom_input.focus()
except NoMatches:
pass
@@ -0,0 +1,48 @@
# Copyright (c) Microsoft. All rights reserved.
"""Agent mode and help text display widget."""
from __future__ import annotations
from rich.text import Text
from textual.reactive import reactive
from textual.widgets import Static
class AgentModeAndHelp(Static):
"""Widget displaying the current agent mode and help text.
Shows the current agent mode (e.g., "plan", "execute") in a colored label,
followed by available commands and help text in a dimmed style. Used in
the fixed bottom area of the console.
Attributes:
mode: Current mode name (e.g., "plan", "execute"), or None if no mode.
mode_color: Rich color string for the mode label (e.g., "yellow", "green").
help_text: Help text to display (e.g., "/exit to quit, /mode to switch").
"""
mode: reactive[str | None] = reactive(None)
mode_color: reactive[str] = reactive("yellow")
help_text: reactive[str] = reactive("")
def render(self) -> Text:
"""Render the mode indicator and help text.
Returns:
Rich Text object with styled mode and help display.
"""
result = Text()
if self.mode:
result.append(f"[{self.mode}]", style=self.mode_color)
if self.help_text:
if self.mode:
result.append(" ")
result.append(self.help_text, style="dim")
if not result.plain:
result.append(" ")
return result
@@ -0,0 +1,31 @@
# Copyright (c) Microsoft. All rights reserved.
"""Mode-colored horizontal rule."""
from __future__ import annotations
from textual.reactive import reactive
from textual.widgets import Static
class PromptRule(Static):
"""A full-width horizontal rule colored by the current agent mode.
Renders a line of '' characters across the terminal width,
colored to match the current mode (e.g., cyan for plan, green for execute).
Attributes:
rule_color: Rich color string for the rule (e.g., "cyan", "green").
"""
rule_color: reactive[str] = reactive("cyan")
def render(self) -> str:
"""Render the horizontal rule.
Returns:
Formatted string with Rich markup.
"""
color = self.rule_color
width = self.size.width or 80
return f"[{color}]{'' * width}[/{color}]"
@@ -0,0 +1,127 @@
# Copyright (c) Microsoft. All rights reserved.
"""Scrolling panel for conversation history display."""
from __future__ import annotations
from typing import TYPE_CHECKING
from textual.widgets import RichLog
if TYPE_CHECKING:
from ..app_state import OutputEntry
class HarnessScrollPanel(RichLog):
"""Scrolling panel for displaying conversation history.
Uses Textual's RichLog widget for efficient append-only rendering with
Rich text formatting support. Automatically scrolls to the bottom when
new entries are added.
For streaming text, the panel uses a truncate-and-rewrite strategy: it
tracks where streaming began in the RichLog lines list, and on each update
truncates back to that point and rewrites the full accumulated text as a
single write. This ensures consistent rendering without line-break artifacts
between streamed chunks.
"""
def __init__(self, **kwargs) -> None:
"""Initialize the scroll panel.
Args:
**kwargs: Additional arguments passed to RichLog.
"""
super().__init__(
**kwargs,
auto_scroll=True, # Automatically scroll to bottom
wrap=True, # Wrap long lines instead of horizontal scroll
markup=True, # Enable Rich markup
highlight=True, # Enable syntax highlighting
)
self._entries: list[OutputEntry] = []
self._is_streaming = False
self._streaming_line_start: int = 0
def append_entry(self, entry: OutputEntry) -> None:
"""Append a new output entry to the conversation history.
Args:
entry: The output entry to append.
"""
self._entries.append(entry)
text = self._format_entry(entry)
self.write(text)
def set_streaming_entry(self, entry: OutputEntry) -> None:
"""Set or update the current streaming entry.
On each update, truncates the RichLog back to where streaming
started, then rewrites the full streaming text as a single block.
This ensures no spurious line breaks between chunks while avoiding
a full rewrite of all entries.
Args:
entry: The streaming entry (will be mutated externally).
"""
if not self._is_streaming:
# First streaming chunk — record where streaming lines begin
self._is_streaming = True
self._entries.append(entry)
self._streaming_line_start = len(self.lines)
# Truncate lines back to where streaming started
if len(self.lines) > self._streaming_line_start:
del self.lines[self._streaming_line_start:]
from textual.geometry import Size
self.virtual_size = Size(self._widest_line_width, len(self.lines))
# Write full streaming text as a single renderable
formatted = self._format_text(entry.text, entry.color)
self.write(formatted)
def end_streaming(self) -> None:
"""End the current streaming mode."""
if self._is_streaming:
self._is_streaming = False
self._streaming_line_start = 0
def _rewrite_all(self) -> None:
"""Clear and rewrite all entries from scratch."""
self.clear()
for entry in self._entries:
self.write(self._format_entry(entry))
def _format_entry(self, entry: OutputEntry) -> str:
"""Format an output entry with Rich markup.
Args:
entry: The entry to format.
Returns:
Formatted string with Rich markup for color and styling.
"""
return self._format_text(entry.text, entry.color)
@staticmethod
def _format_text(text: str, color: str | None) -> str:
"""Format text with optional Rich color markup.
Args:
text: The text to format.
color: Optional Rich color name.
Returns:
Formatted string.
"""
if color:
return f"[{color}]{text}[/{color}]"
return text
def clear_history(self) -> None:
"""Clear all conversation history from the panel."""
self._entries.clear()
self._is_streaming = False
self._streaming_line_start = 0
self.clear()
@@ -0,0 +1,102 @@
# Copyright (c) Microsoft. All rights reserved.
"""Text input widget with inline prompt for the harness console."""
from __future__ import annotations
from textual import on
from textual.app import ComposeResult
from textual.containers import Horizontal
from textual.message import Message
from textual.reactive import reactive
from textual.widget import Widget
from textual.widgets import Input, Label
class HarnessTextInput(Widget):
"""Text input widget with a prompt label on the left.
Displays a prompt (e.g., "> ") followed by a borderless input field.
Sits between the two mode-colored horizontal rules.
Attributes:
prompt: The prompt text displayed on the left (e.g., "> ").
placeholder: Placeholder text shown when the input is empty.
"""
prompt: reactive[str] = reactive("> ")
placeholder: reactive[str] = reactive("")
class Submitted(Message):
"""Message sent when the input is submitted.
Attributes:
value: The submitted text value.
"""
def __init__(self, value: str) -> None:
"""Initialize the Submitted message.
Args:
value: The submitted text value.
"""
self.value = value
super().__init__()
def compose(self) -> ComposeResult:
"""Compose the prompt label and input field.
Yields:
A horizontal container with the prompt and input field.
"""
with Horizontal(classes="prompt-container"):
yield Label(self.prompt, classes="prompt-label", id="prompt-label")
yield Input(placeholder=self.placeholder, classes="input-field", id="input-field")
def watch_prompt(self, new_prompt: str) -> None:
"""Update the prompt label when the prompt attribute changes.
Args:
new_prompt: The new prompt text.
"""
try:
label = self.query_one("#prompt-label", Label)
label.update(new_prompt)
except Exception:
pass
def watch_placeholder(self, new_placeholder: str) -> None:
"""Update the input placeholder when the placeholder attribute changes.
Args:
new_placeholder: The new placeholder text.
"""
try:
input_field = self.query_one("#input-field", Input)
input_field.placeholder = new_placeholder
except Exception:
# Input doesn't exist yet (before compose), ignore
pass
@on(Input.Submitted)
def on_input_submitted(self, event: Input.Submitted) -> None:
"""Handle input submission.
Clears the input field and posts a Submitted message with the value.
Args:
event: The Input.Submitted event.
"""
value = event.value
event.input.clear()
self.post_message(self.Submitted(value))
def focus_input(self) -> None:
"""Focus the input field."""
input_field = self.query_one(".input-field", Input)
input_field.focus()
def clear_input(self) -> None:
"""Clear the input field."""
input_field = self.query_one(".input-field", Input)
input_field.clear()
@@ -0,0 +1,503 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tool call formatters for displaying function calls in the harness console.
This module provides formatters that convert raw function call content into
human-readable display strings. Each formatter handles specific tool patterns
(e.g., web_search, todos_*, etc.) and the FallbackToolFormatter provides
generic formatting for any unmatched tools.
Usage:
from harness.console.formatters import build_default_formatters, format_tool_call
from agent_framework import Content
call = Content.from_function_call(
call_id="call_1",
name="web_search",
arguments={"query": "Python async"}
)
formatters = build_default_formatters()
result = format_tool_call(formatters, call) # "web_search (Python async)"
"""
from __future__ import annotations
import contextlib
import json
from abc import ABC, abstractmethod
from typing import Any
from agent_framework import Content
# region Helper Functions
def get_argument_value(call: Content, param_name: str) -> Any:
"""Extract an argument value from a function call.
Handles both dict and JSON string arguments.
Args:
call: The function call content.
param_name: The parameter name to extract.
Returns:
The argument value, or None if not found.
"""
if call.arguments is None:
return None
if isinstance(call.arguments, str):
# arguments is a JSON string, parse it
try:
args_dict = json.loads(call.arguments)
except (json.JSONDecodeError, TypeError):
return None
if not isinstance(args_dict, dict):
return None
elif isinstance(call.arguments, dict):
args_dict = call.arguments
else:
return None
return args_dict.get(param_name)
def as_int_list(value: Any) -> list[int] | None:
"""Convert a value to a list of integers, or None if not possible.
Args:
value: The value to convert (should be a list).
Returns:
A list of integers, or None if conversion fails.
"""
if not isinstance(value, list):
return None
result: list[int] = []
for item in value:
if isinstance(item, int):
result.append(item)
else:
with contextlib.suppress(ValueError, TypeError):
result.append(int(item))
return result if result else None
def as_dict_list(value: Any) -> list[dict[str, Any]] | None:
"""Convert a value to a list of dicts, or None if not possible.
Args:
value: The value to convert (should be a list).
Returns:
A list of dicts, or None if value is not a list of dicts.
"""
if not isinstance(value, list):
return None
result: list[dict[str, Any]] = []
for item in value:
if isinstance(item, dict):
result.append(item)
return result if result else None
def truncate(text: str, max_length: int) -> str:
"""Truncate a string to the specified maximum length, appending an ellipsis if truncated.
Args:
text: The text to truncate.
max_length: The maximum length.
Returns:
The truncated string.
"""
return text if len(text) <= max_length else text[:max_length] + ""
# endregion
# region Base Class
class ToolCallFormatter(ABC):
"""Base class for tool call formatters that produce human-readable display strings
for function call content items shown in the console.
"""
@abstractmethod
def can_format(self, call: Content) -> bool:
"""Return True if this formatter can handle the given function call.
Args:
call: The function call content to check.
Returns:
True if this formatter should be used; otherwise False.
"""
...
@abstractmethod
def format_detail(self, call: Content) -> str | None:
"""Return the detail portion of the formatted output for the given tool call,
or None if only the tool name should be displayed.
Args:
call: The function call content to format.
Returns:
A detail string to append after the tool name, or None.
"""
...
# endregion
# region Concrete Formatters
class FallbackToolFormatter(ToolCallFormatter):
"""Catch-all formatter that handles any tool not matched by a more specific formatter.
Displays a generic summary of the tool's arguments. This formatter should always be
placed last in the formatter list.
"""
def can_format(self, call: Content) -> bool:
"""Always returns True - this formatter matches everything."""
return True
def format_detail(self, call: Content) -> str | None:
"""Format arguments as generic (key: value, ...) pairs."""
if call.arguments is None:
return None
# Parse arguments
if isinstance(call.arguments, str):
try:
args_dict = json.loads(call.arguments)
except (json.JSONDecodeError, TypeError):
return None
elif isinstance(call.arguments, dict):
args_dict = call.arguments
else:
return None
if not args_dict:
return None
# Build argument list
parts: list[str] = []
for key, value in args_dict.items():
if value is None:
continue
# Convert value to string
if isinstance(value, bool):
str_value = "true" if value else "false"
elif isinstance(value, (int, float)):
str_value = str(value)
elif isinstance(value, str):
str_value = value
else:
# Complex types - skip for now
continue
parts.append(f"{key}: {truncate(str_value, 40)}")
return f"({', '.join(parts)})" if parts else None
class WebSearchToolFormatter(ToolCallFormatter):
"""Formats web_search tool calls, showing the search query."""
def can_format(self, call: Content) -> bool:
"""Match web_search tool calls."""
return call.name == "web_search"
def format_detail(self, call: Content) -> str | None:
"""Extract and format the query parameter."""
value = get_argument_value(call, "query")
return f"({value})" if value else None
class TodoToolFormatter(ToolCallFormatter):
"""Formats todos_* tool calls with tree-view output for added items
and structured output for complete/remove operations.
"""
def can_format(self, call: Content) -> bool:
"""Match todos_* tool calls."""
return call.name is not None and call.name.startswith("todos_")
def format_detail(self, call: Content) -> str | None:
"""Format based on the specific todos operation."""
if call.name == "todos_add":
return self._format_add_todos(call)
if call.name == "todos_complete":
return self._format_complete_todos(call)
if call.name == "todos_remove":
return self._format_id_list(call, "ids", "Remove")
return None
def _format_add_todos(self, call: Content) -> str | None:
"""Format todos_add with tree view of titles."""
todos = as_dict_list(get_argument_value(call, "todos"))
if not todos:
return None
titles: list[str] = []
for todo in todos:
title = todo.get("title")
if title and isinstance(title, str):
titles.append(title)
if not titles:
return None
# Build tree view
count = len(titles)
plural = "s" if count != 1 else ""
lines = [f"({count} item{plural})"]
for i, title in enumerate(titles):
connector = "├─" if i < count - 1 else "└─"
lines.append(f"\n {connector} {title}")
return "".join(lines)
def _format_complete_todos(self, call: Content) -> str | None:
"""Format todos_complete with tree view of IDs and reasons."""
items = as_dict_list(get_argument_value(call, "items"))
if not items:
return None
entries: list[tuple[int, str | None]] = []
for item in items:
todo_id = item.get("id")
if not isinstance(todo_id, int):
continue
reason = item.get("reason")
reason_str = str(reason) if reason is not None and not isinstance(reason, str) else reason
entries.append((todo_id, reason_str))
if not entries:
return None
# Build tree view
lines: list[str] = []
for i, (todo_id, reason) in enumerate(entries):
connector = "├─" if i < len(entries) - 1 else "└─"
line = f"\n {connector} Complete #{todo_id}"
if reason:
line += f"{truncate(reason, 80)}"
lines.append(line)
return "".join(lines)
def _format_id_list(self, call: Content, param_name: str, verb: str) -> str | None:
"""Format a list of IDs with a verb (e.g., Remove #1, Remove #2)."""
ids = as_int_list(get_argument_value(call, param_name))
if not ids:
return None
lines: list[str] = []
for i, todo_id in enumerate(ids):
connector = "├─" if i < len(ids) - 1 else "└─"
lines.append(f"\n {connector} {verb} #{todo_id}")
return "".join(lines)
class ModeToolFormatter(ToolCallFormatter):
"""Formats AgentMode_* tool calls, showing the target mode for Set operations."""
def can_format(self, call: Content) -> bool:
"""Match AgentMode_* tool calls."""
return call.name is not None and call.name.startswith("AgentMode_")
def format_detail(self, call: Content) -> str | None:
"""Format based on the specific AgentMode operation."""
if call.name == "AgentMode_Set":
value = get_argument_value(call, "mode")
return f"({value})" if value else None
return None
class BackgroundAgentToolFormatter(ToolCallFormatter):
"""Formats BackgroundAgents_* tool calls with human-readable details
for task start, continue, wait, and result retrieval operations.
"""
def can_format(self, call: Content) -> bool:
"""Match BackgroundAgents_* tool calls."""
return call.name is not None and call.name.startswith("BackgroundAgents_")
def format_detail(self, call: Content) -> str | None:
"""Format based on the specific BackgroundAgents operation."""
if call.name == "BackgroundAgents_StartTask":
return self._format_start_background_task(call)
if call.name == "BackgroundAgents_WaitForFirstCompletion":
return self._format_id_list(call, "taskIds", "Wait for")
if call.name == "BackgroundAgents_GetTaskResults":
return self._format_single_id(call, "taskId")
if call.name == "BackgroundAgents_ContinueTask":
return self._format_continue_task(call)
if call.name == "BackgroundAgents_ClearCompletedTask":
return self._format_single_id(call, "taskId")
return None
def _format_start_background_task(self, call: Content) -> str | None:
"""Format StartTask with agent name and description."""
agent_name = get_argument_value(call, "agentName")
description = get_argument_value(call, "description")
if agent_name is None and description is None:
return None
lines: list[str] = []
if agent_name is not None and description is not None:
lines.append(f"\n ├─ Agent: {agent_name}")
lines.append(f'\n └─ "{truncate(description, 80)}"')
elif agent_name is not None:
lines.append(f"\n └─ Agent: {agent_name}")
else:
lines.append(f'\n └─ "{truncate(description, 80)}"') # type: ignore[arg-type]
return "".join(lines)
def _format_id_list(self, call: Content, param_name: str, verb: str) -> str | None:
"""Format a list of task IDs with a verb."""
ids = as_int_list(get_argument_value(call, param_name))
if not ids:
return None
lines: list[str] = []
for i, task_id in enumerate(ids):
connector = "├─" if i < len(ids) - 1 else "└─"
lines.append(f"\n {connector} {verb} #{task_id}")
return "".join(lines)
def _format_single_id(self, call: Content, param_name: str) -> str | None:
"""Format a single task ID in parentheses."""
task_id = get_argument_value(call, param_name)
if isinstance(task_id, int):
return f"(task #{task_id})"
return None
def _format_continue_task(self, call: Content) -> str | None:
"""Format ContinueTask with task ID and optional text."""
task_id = get_argument_value(call, "taskId")
text = get_argument_value(call, "text")
if not isinstance(task_id, int):
return None
if text:
lines = [
f"\n ├─ Task #{task_id}",
f'\n └─ "{truncate(text, 80)}"',
]
return "".join(lines)
return f"\n └─ Task #{task_id}"
class FileMemoryToolFormatter(ToolCallFormatter):
"""Formats FileMemory_* tool calls, showing file names and search patterns
with tree-view corners for save operations.
"""
def can_format(self, call: Content) -> bool:
"""Match FileMemory_* tool calls."""
return call.name is not None and call.name.startswith("FileMemory_")
def format_detail(self, call: Content) -> str | None:
"""Format based on the specific FileMemory operation."""
if call.name == "FileMemory_SaveFile":
return self._format_save_file(call)
if call.name in ("FileMemory_ReadFile", "FileMemory_DeleteFile"):
value = get_argument_value(call, "fileName")
return f"({value})" if value else None
if call.name == "FileMemory_SearchFiles":
return self._format_search_files(call)
return None
def _format_save_file(self, call: Content) -> str | None:
"""Format SaveFile with file name and description indicator."""
file_name = get_argument_value(call, "fileName")
description = get_argument_value(call, "description")
if not file_name:
return None
if description:
return f"\n └─ {file_name} (with description)"
return f"\n └─ {file_name}"
def _format_search_files(self, call: Content) -> str | None:
"""Format SearchFiles with regex pattern and optional file pattern."""
pattern = get_argument_value(call, "regexPattern")
file_pattern = get_argument_value(call, "filePattern")
if not pattern:
return None
if file_pattern:
return f"(/{pattern}/ in {file_pattern})"
return f"(/{pattern}/)"
# endregion
# region Public API Functions
def format_tool_call(formatters: list[ToolCallFormatter], call: Content) -> str:
"""Format a tool call using the first matching formatter from the provided list.
Returns "{toolName} {detail}" when a formatter produces detail,
or just "{toolName}" otherwise.
Args:
formatters: List of formatters to try in order.
call: The function call content to format.
Returns:
Formatted string representation of the tool call.
"""
for formatter in formatters:
if formatter.can_format(call):
detail = formatter.format_detail(call)
tool_name = call.name or "Unknown"
return f"{tool_name} {detail}" if detail is not None else tool_name
return call.name or "Unknown"
def build_default_formatters() -> list[ToolCallFormatter]:
"""Create the default list of tool call formatters.
The FallbackToolFormatter is always last. Users can call this function
and combine the result with their own formatters.
Returns:
A list of all built-in tool call formatters.
"""
return [
TodoToolFormatter(),
ModeToolFormatter(),
BackgroundAgentToolFormatter(),
FileMemoryToolFormatter(),
WebSearchToolFormatter(),
FallbackToolFormatter(),
]
# endregion
@@ -0,0 +1,87 @@
# Copyright (c) Microsoft. All rights reserved.
"""Main entry point for the harness console.
Provides the top-level run_agent_async() function that creates and runs
the Textual-based harness console application.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from .app import HarnessApp
from .observers import build_default_observers
if TYPE_CHECKING:
from agent_framework import Agent, AgentSession
from .commands import CommandHandler
from .observers.base import ConsoleObserver
async def run_agent_async(
agent: Agent,
*,
session: AgentSession | None = None,
observers: list[ConsoleObserver] | None = None,
command_handlers: list[CommandHandler] | None = None,
mode_colors: dict[str, str] | None = None,
initial_mode: str | None = None,
placeholder: str = "Type a message and press Enter...",
title: str = "Harness Console",
max_context_window_tokens: int | None = None,
max_output_tokens: int | None = None,
) -> None:
"""Run the harness console with the given agent.
This is the main entry point for the harness console. Creates a Textual
application with the configured observers and runs it until the user exits.
Args:
agent: The agent to run conversations with.
session: Optional agent session for conversation history.
observers: List of console observers. If None, uses defaults.
command_handlers: List of command handlers. If None, auto-detected from agent.
mode_colors: Mapping of mode names to Rich color strings.
initial_mode: Initial agent mode text.
placeholder: Input placeholder text.
title: Application title.
max_context_window_tokens: Optional max context window size for usage display.
max_output_tokens: Optional max output tokens for usage display.
Example:
.. code-block:: python
from agent_framework import Agent
from agent_framework.openai import OpenAIChatClient
from console import run_agent_async
agent = Agent(
client=OpenAIChatClient(),
instructions="You are helpful.",
)
await run_agent_async(agent)
"""
resolved_observers = observers or build_default_observers()
resolved_mode_colors = mode_colors or {
"plan": "cyan",
"execute": "green",
}
resolved_session = session or agent.create_session()
app = HarnessApp(
agent=agent,
observers=resolved_observers,
session=resolved_session,
mode_colors=resolved_mode_colors,
initial_mode=initial_mode,
placeholder=placeholder,
title=title,
max_context_window_tokens=max_context_window_tokens,
max_output_tokens=max_output_tokens,
command_handlers=command_handlers,
)
await app.run_async()
@@ -0,0 +1,122 @@
# Copyright (c) Microsoft. All rights reserved.
"""Console observers for agent streaming lifecycle.
This module provides observers that display events during agent streaming
and collect follow-up actions. All observers use the IUXStateDriver interface
to update the UI.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
from .base import ConsoleObserver
from .error_display import ErrorDisplayObserver
from .planning_output import PlanningOutputObserver
from .reasoning_display import ReasoningDisplayObserver
from .text_output import TextOutputObserver
from .tool_approval import ToolApprovalObserver
from .tool_call_display import ToolCallDisplayObserver
from .usage_display import UsageDisplayObserver
if TYPE_CHECKING:
from agent_framework import Agent
def build_default_observers() -> list[ConsoleObserver]:
"""Build the default set of observers for the harness console.
Returns a standard observer list covering:
- Text output (streaming text display)
- Tool call display (formatted tool invocations)
- Error display (error messages)
- Usage display (token counts)
- Reasoning display (reasoning/thinking blocks)
- Tool approval (user approval for tool calls)
Note: PlanningOutputObserver is NOT included here because it requires
a mode_provider. Use build_observers_with_planning() for agents that
have an AgentModeProvider (i.e. agents created with create_harness_agent).
Returns:
List of default console observers.
"""
return [
TextOutputObserver(),
ToolCallDisplayObserver(),
ErrorDisplayObserver(),
UsageDisplayObserver(),
ReasoningDisplayObserver(),
ToolApprovalObserver(),
]
def build_observers_with_planning(
agent: Agent,
plan_mode_name: str = "plan",
execution_mode_name: str = "execute",
*,
mode_colors: dict[str, str] | None = None,
) -> list[ConsoleObserver]:
"""Build observers with planning support (structured output in plan mode).
Replaces TextOutputObserver with PlanningOutputObserver, which configures
structured JSON output via response_format when in plan mode. This enables
the list picker UI for clarification and approval questions.
Requires that the agent has an AgentModeProvider in its context_providers
(automatically added by create_harness_agent).
Args:
agent: The agent to resolve the AgentModeProvider from.
plan_mode_name: The mode name that represents planning mode.
execution_mode_name: The mode name to switch to on approval.
mode_colors: Optional mapping of mode names to Rich color strings.
Returns:
List of observers with planning support.
Raises:
ValueError: If the agent has no AgentModeProvider.
"""
from agent_framework import AgentModeProvider
mode_provider = next(
(p for p in agent.context_providers if isinstance(p, AgentModeProvider)),
None,
)
if mode_provider is None:
msg = (
"Planning observers require an AgentModeProvider on the agent. "
"Use create_harness_agent() or add AgentModeProvider to context_providers."
)
raise ValueError(msg)
return [
ToolCallDisplayObserver(),
ToolApprovalObserver(),
ErrorDisplayObserver(),
ReasoningDisplayObserver(),
UsageDisplayObserver(),
PlanningOutputObserver(
mode_provider,
plan_mode_name,
execution_mode_name,
mode_colors=mode_colors,
),
]
__all__ = [
"ConsoleObserver",
"ErrorDisplayObserver",
"PlanningOutputObserver",
"ReasoningDisplayObserver",
"TextOutputObserver",
"ToolApprovalObserver",
"ToolCallDisplayObserver",
"UsageDisplayObserver",
"build_default_observers",
"build_observers_with_planning",
]
@@ -0,0 +1,125 @@
# Copyright (c) Microsoft. All rights reserved.
"""Base class for console observers.
Observers participate in the agent streaming lifecycle, displaying events
and optionally returning follow-up actions.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
if TYPE_CHECKING:
from agent_framework import Agent, Content, Message
from ..app_state import FollowUpAction
from ..state_driver import IUXStateDriver
class ConsoleObserver:
"""Base class for console observers.
Observers participate in the agent streaming lifecycle, displaying
events (tool calls, errors, reasoning, etc.) and optionally returning
follow-up actions (questions, approval requests).
All methods have default no-op implementations, so subclasses only
override the methods they need.
"""
def configure_run_options(
self,
options: dict[str, Any],
agent: Agent,
session: Any,
) -> None:
"""Configure run options before agent invocation.
Override to set options such as response_format, max_tokens, etc.
Args:
options: Dictionary of chat options to modify.
agent: The AI agent.
session: The agent session.
"""
pass
async def on_response_update(
self,
ux: IUXStateDriver,
update: Message,
agent: Agent,
session: Any,
) -> None:
"""Called for each response update chunk.
Override to inspect update-level metadata or handle provider-specific
events in the raw representation.
Args:
ux: The UX state driver for UI updates.
update: The message update chunk.
agent: The AI agent.
session: The agent session.
"""
pass
async def on_content(
self,
ux: IUXStateDriver,
content: Content,
agent: Agent,
session: Any,
) -> None:
"""Called for each content item in the response.
Override to handle specific content types (function calls, errors, etc.).
Args:
ux: The UX state driver for UI updates.
content: The content item from the response.
agent: The AI agent.
session: The agent session.
"""
pass
async def on_text(
self,
ux: IUXStateDriver,
text: str,
agent: Agent,
session: Any,
) -> None:
"""Called for each text chunk in the response.
Override to accumulate and display streaming text.
Args:
ux: The UX state driver for UI updates.
text: The text chunk.
agent: The AI agent.
session: The agent session.
"""
pass
async def on_stream_complete(
self,
ux: IUXStateDriver,
agent: Agent,
session: Any,
) -> list[FollowUpAction] | None:
"""Called when streaming completes.
Override to return follow-up actions (questions to ask the user,
messages to inject into the next turn, etc.).
Args:
ux: The UX state driver for UI updates.
agent: The AI agent.
session: The agent session.
Returns:
Optional list of follow-up actions to queue, or None.
"""
return None
@@ -0,0 +1,72 @@
# Copyright (c) Microsoft. All rights reserved.
"""Error display observer for showing errors."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from .base import ConsoleObserver
if TYPE_CHECKING:
from agent_framework import Agent, Content
from ..state_driver import IUXStateDriver
class ErrorDisplayObserver(ConsoleObserver):
"""Displays error content from the agent response.
Shows errors with an prefix in red to make them easily visible.
"""
async def on_content(
self,
ux: IUXStateDriver,
content: Content,
agent: Agent,
session: Any,
) -> None:
"""Display error content.
Args:
ux: The UX state driver for UI updates.
content: The content item to check for errors.
agent: The AI agent.
session: The agent session.
"""
# Check if this is an error content type
# The exact content type check depends on the agent framework's Content class
if hasattr(content, "type") and content.type == "error":
error_text = self._format_error(content)
ux.append_info_line(error_text, "red")
elif getattr(content, "error", None):
error_text = f"❌ Error: {content.error}" # type: ignore[reportAttributeAccessIssue]
ux.append_info_line(error_text, "red")
def _format_error(self, content: Content) -> str:
"""Format error content for display.
Args:
content: The error content.
Returns:
Formatted error string.
"""
error_text = "❌ Error"
# Try to extract error message
if hasattr(content, "message"):
error_text += f": {content.message}"
elif hasattr(content, "text"):
error_text += f": {content.text}"
# Try to add error code if available
if hasattr(content, "error_code") and content.error_code:
error_text += f" (code: {content.error_code})"
# Try to add details if available
if hasattr(content, "details") and getattr(content, "details", None):
error_text += f"{content.details}" # type: ignore[reportAttributeAccessIssue]
return error_text
@@ -0,0 +1,71 @@
# Copyright (c) Microsoft. All rights reserved.
"""Pydantic models for structured planning output.
These models define the JSON schema that the agent produces when in planning
mode via `response_format`. The schema enables consistent rendering of
clarification questions and approval requests in the console UI.
"""
from __future__ import annotations
from enum import Enum
from pydantic import BaseModel, Field
class PlanningResponseType(str, Enum):
"""Type of planning response from the agent."""
CLARIFICATION = "clarification"
"""The agent needs clarification and presents options for the user to choose from."""
APPROVAL = "approval"
"""The agent is seeking approval to proceed with execution."""
class PlanningQuestion(BaseModel):
"""A single question or item within a PlanningResponse.
For clarification: contains the question text and optional choices.
For approval: contains the plan summary for the user to approve.
"""
message: str = Field(
description=(
"For clarifications, this has the question that needs to be clarified "
"with the user. For approvals, this would contain a summary of the "
"execution plan that the user needs to approve."
),
)
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."
),
)
class PlanningResponse(BaseModel):
"""Structured response from the agent while in planning mode.
Used with structured output (`response_format`) to enable consistent
rendering of clarification questions and approval requests.
"""
type: PlanningResponseType = Field(
description=(
"Use 'clarification' when you need clarification around the user "
"request and you want to present the user with options to choose from. "
"Use 'approval' when you are ready to start execution, but need "
"approval to start executing."
),
)
questions: list[PlanningQuestion] = Field(
description=(
"For clarifications, this has one or more questions to ask the user "
"(each with choices). For approvals, this has exactly one item "
"containing the plan summary for the user to approve."
),
)
@@ -0,0 +1,242 @@
# Copyright (c) Microsoft. All rights reserved.
"""Planning output observer for structured agent responses in plan mode.
In planning mode, this observer configures structured JSON output via
response_format, collects streamed text silently, then deserializes the
result as a PlanningResponse to present clarification/approval questions.
In execution mode, text is streamed through directly.
"""
from __future__ import annotations
import json
from typing import TYPE_CHECKING, Any
from rich.markup import escape
from ..app_state import (
ChoiceFollowUpQuestion,
FollowUpAction,
TextFollowUpQuestion,
)
from .base import ConsoleObserver
from .planning_models import PlanningResponse, PlanningResponseType
if TYPE_CHECKING:
from agent_framework import Agent, AgentModeProvider, Message
from ..state_driver import IUXStateDriver
class PlanningOutputObserver(ConsoleObserver):
"""Mode-aware observer that uses structured output in plan mode.
In planning mode:
- Configures response_format to PlanningResponse schema
- Collects streamed text silently
- Deserializes JSON into PlanningResponse
- Builds follow-up questions (clarification or approval)
In execution mode:
- Streams text directly to the UX driver
If JSON parsing fails, falls back to rendering the raw text as regular
output so the user always sees what the agent produced.
"""
def __init__(
self,
mode_provider: AgentModeProvider,
plan_mode_name: str,
execution_mode_name: str,
*,
mode_colors: dict[str, str] | None = None,
) -> None:
"""Initialize the planning output observer.
Args:
mode_provider: The mode provider for reading/switching modes.
plan_mode_name: The mode name that represents planning mode.
execution_mode_name: The mode name to switch to on approval.
mode_colors: Optional mapping of mode names to Rich color strings.
"""
self._mode_provider = mode_provider
self._plan_mode_name = plan_mode_name
self._execution_mode_name = execution_mode_name
self._mode_colors = mode_colors or {}
self._text_collector: list[str] = []
def configure_run_options(
self,
options: dict[str, Any],
agent: Agent,
session: Any,
) -> None:
"""Set response_format to PlanningResponse when in plan mode."""
if self._is_planning_mode(session):
options["response_format"] = PlanningResponse
async def on_text(
self,
ux: IUXStateDriver,
text: str,
agent: Agent,
session: Any,
) -> None:
"""Collect text in plan mode; stream through in execute mode."""
if self._is_planning_mode_from_ux(ux):
self._text_collector.append(text)
else:
ux.write_text(escape(text))
async def on_stream_complete(
self,
ux: IUXStateDriver,
agent: Agent,
session: Any,
) -> list[FollowUpAction] | None:
"""Parse collected text as PlanningResponse and build follow-up actions."""
if not self._is_planning_mode_from_ux(ux):
self._text_collector.clear()
return None
collected_text = "".join(self._text_collector)
self._text_collector.clear()
if not collected_text.strip():
return None
# Attempt to deserialize structured response
try:
planning_response = PlanningResponse.model_validate_json(collected_text)
except (json.JSONDecodeError, ValueError):
# JSON parsing failed — fall back to rendering as regular text
ux.write_text(escape(collected_text))
return None
if planning_response.type == PlanningResponseType.CLARIFICATION:
return self._build_clarification_actions(planning_response)
if planning_response.type == PlanningResponseType.APPROVAL:
if not planning_response.questions:
ux.append_info_line("(approval response had no content)", "yellow")
return None
question = planning_response.questions[0]
return [self._build_approval_action(question, session)]
# Unexpected type — fall back to rendering as regular text
ux.write_text(escape(collected_text))
return None
def _is_planning_mode(self, session: Any) -> bool:
"""Check if session is in planning mode."""
from agent_framework import get_agent_mode
try:
current_mode = get_agent_mode(session)
except (AttributeError, TypeError):
return True # No mode provider → treat as planning
return current_mode.lower() == self._plan_mode_name.lower()
def _is_planning_mode_from_ux(self, ux: IUXStateDriver) -> bool:
"""Check if UX is in planning mode."""
current = ux.current_mode
if current is None:
return True
return current.lower() == self._plan_mode_name.lower()
def _build_clarification_actions(
self,
response: PlanningResponse,
) -> list[FollowUpAction]:
"""Build follow-up questions for clarification."""
actions: list[FollowUpAction] = []
for question in response.questions:
prompt = question.message
cont = self._make_clarification_continuation(prompt)
if question.choices and len(question.choices) > 0:
actions.append(
ChoiceFollowUpQuestion(
prompt=prompt,
choices=question.choices,
allow_custom_text=True,
continuation=cont,
)
)
else:
actions.append(
TextFollowUpQuestion(
prompt=prompt,
continuation=cont,
)
)
return actions
@staticmethod
def _make_clarification_continuation(prompt: str):
"""Create a clarification continuation closure capturing the prompt."""
async def continuation(
answer: str,
ux: IUXStateDriver,
) -> Message | None:
if not answer.strip():
ux.append_info_line(f"🔹 {prompt}\n └─ (no answer)", "dim")
return None
ux.append_info_line(f"🔹 {prompt}\n └─ [green]{answer}[/green]", "dim")
from agent_framework import Message
return Message(role="user", contents=[f"Q: {prompt}\nA: {answer}"])
return continuation
def _build_approval_action(
self,
question: Any,
session: Any,
) -> ChoiceFollowUpQuestion:
"""Build the approval follow-up question."""
approve_option = "Approve and switch to execute mode"
prompt = question.message
async def continuation(
selection: str,
ux: IUXStateDriver,
) -> Message | None:
ux.append_info_line(
f"🔹 {prompt}\n └─ [green]{selection}[/green]",
"dim",
)
if selection == approve_option:
from agent_framework import set_agent_mode
set_agent_mode(session, self._execution_mode_name)
exec_color = self._mode_colors.get(self._execution_mode_name)
ux.set_mode(self._execution_mode_name, exec_color)
ux.append_info_line(
f"✅ Switched to {self._execution_mode_name} mode.",
exec_color,
)
from agent_framework import Message
return Message(role="user", contents=["Approved"])
# Custom freeform input — treat as suggested changes
from agent_framework import Message
return Message(role="user", contents=[selection])
return ChoiceFollowUpQuestion(
prompt=prompt,
choices=[approve_option],
allow_custom_text=True,
continuation=continuation,
)
@@ -0,0 +1,80 @@
# Copyright (c) Microsoft. All rights reserved.
"""Reasoning display observer for showing thinking content."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from rich.markup import escape
from .base import ConsoleObserver
if TYPE_CHECKING:
from agent_framework import Agent, Content
from ..state_driver import IUXStateDriver
class ReasoningDisplayObserver(ConsoleObserver):
"""Displays reasoning/thinking content from the agent.
Some models (like o1) provide reasoning steps that show their
internal thought process. This observer displays them with a 💭 prefix
in a dimmed style.
"""
async def on_content(
self,
ux: IUXStateDriver,
content: Content,
agent: Agent,
session: Any,
) -> None:
"""Display reasoning content.
Args:
ux: The UX state driver for UI updates.
content: The content item to check for reasoning.
agent: The AI agent.
session: The agent session.
"""
reasoning_text = self._extract_reasoning(content)
if reasoning_text:
# Display reasoning in dim style to differentiate from main output
ux.append_info_line(f"💭 {escape(reasoning_text)}", "dim")
def _extract_reasoning(self, content: Content) -> str | None:
"""Extract reasoning text from content.
Args:
content: The content item to extract reasoning from.
Returns:
The reasoning text, or None if no reasoning is present.
"""
# Check for reasoning content type
if hasattr(content, "type") and content.type in {"text_reasoning", "reasoning"}:
if hasattr(content, "text"):
return content.text
content_attr = getattr(content, "content", None)
if content_attr:
return str(content_attr)
# Check for reasoning attribute
reasoning = getattr(content, "reasoning", None)
if reasoning is not None:
if isinstance(reasoning, str):
return reasoning
if hasattr(reasoning, "text"):
return reasoning.text
# Check for thinking attribute (alternative name)
thinking = getattr(content, "thinking", None)
if thinking is not None:
if isinstance(thinking, str):
return thinking
if hasattr(thinking, "text"):
return thinking.text
return None
@@ -0,0 +1,59 @@
# Copyright (c) Microsoft. All rights reserved.
"""Text output observer for streaming agent text."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from rich.markup import escape
from .base import ConsoleObserver
if TYPE_CHECKING:
from agent_framework import Agent
from ..state_driver import IUXStateDriver
class TextOutputObserver(ConsoleObserver):
"""Displays streaming text output from the agent.
Writes text chunks incrementally to the UX state driver as they arrive,
allowing real-time display during streaming.
"""
async def on_text(
self,
ux: IUXStateDriver,
text: str,
agent: Agent,
session: Any,
) -> None:
"""Write each text chunk directly to the UX driver.
Args:
ux: The UX state driver for UI updates.
text: The text chunk to display.
agent: The AI agent.
session: The agent session.
"""
ux.write_text(escape(text))
async def on_stream_complete(
self,
ux: IUXStateDriver,
agent: Agent,
session: Any,
) -> list | None:
"""No-op on stream complete (state managed by UX driver).
Args:
ux: The UX state driver for UI updates.
agent: The AI agent.
session: The agent session.
Returns:
None (no follow-up actions).
"""
return None
@@ -0,0 +1,139 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tool approval observer for user confirmation of tool calls.
Detects function_approval_request content items during streaming, displays
approval notifications, and after the stream completes presents one
ChoiceFollowUpQuestion per pending approval request.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from ..app_state import ChoiceFollowUpQuestion, FollowUpAction
from .base import ConsoleObserver
if TYPE_CHECKING:
from agent_framework import Agent, Content, Message
from ..state_driver import IUXStateDriver
class ToolApprovalObserver(ConsoleObserver):
"""Asks user to approve tool calls before execution.
Collects `function_approval_request` content during streaming and presents
a multi-choice approval question for each after the stream completes.
The continuation builds a `function_approval_response` Content to inject
into the next agent turn.
"""
def __init__(self) -> None:
"""Initialize the tool approval observer."""
self._approval_requests: list[Content] = []
async def on_content(
self,
ux: IUXStateDriver,
content: Content,
agent: Agent,
session: Any,
) -> None:
"""Collect function_approval_request content for approval.
Args:
ux: The UX state driver for UI updates.
content: The content item to check.
agent: The AI agent.
session: The agent session.
"""
if content.type == "function_approval_request":
self._approval_requests.append(content)
tool_name = self._format_tool_name(content)
ux.append_info_line(f"⚠️ Approval needed: {tool_name}", "yellow")
async def on_stream_complete(
self,
ux: IUXStateDriver,
agent: Agent,
session: Any,
) -> list[FollowUpAction] | None:
"""Build approval questions for collected requests.
Args:
ux: The UX state driver for UI updates.
agent: The AI agent.
session: The agent session.
Returns:
List of ChoiceFollowUpQuestions, one per approval request.
"""
if not self._approval_requests:
return None
actions: list[FollowUpAction] = []
for request in self._approval_requests:
actions.append(self._build_approval_question(request))
self._approval_requests.clear()
return actions
def _build_approval_question(self, request: Content) -> ChoiceFollowUpQuestion:
"""Build a multi-choice approval question for a single request."""
tool_name = self._format_tool_name(request)
prompt = f"🔐 Tool approval: {tool_name}"
# TODO(westey-m): Add "Always approve" options when the framework supports
# CreateAlwaysApproveToolResponse / CreateAlwaysApproveToolWithArgumentsResponse.
choices = [
"Approve this call",
"Deny",
]
async def continuation(
selection: str,
ux: IUXStateDriver,
) -> Message | None:
from agent_framework import Message
if selection == "Deny":
response_content = request.to_function_approval_response(approved=False)
action_label = "❌ Denied"
color = "red"
else:
response_content = request.to_function_approval_response(approved=True)
action_label = "✅ Approved"
color = "green"
ux.append_info_line(
f"🔹 {prompt}\n └─ [{color}]{action_label}[/{color}]",
"dim",
)
return Message(role="user", contents=[response_content])
return ChoiceFollowUpQuestion(
prompt=prompt,
choices=choices,
allow_custom_text=False,
continuation=continuation,
)
@staticmethod
def _format_tool_name(content: Content) -> str:
"""Extract a readable tool name from approval request content."""
# The function_call is stored on the approval request content
function_call = getattr(content, "function_call", None)
if function_call is not None:
from ..formatters import build_default_formatters, format_tool_call
try:
return format_tool_call(build_default_formatters(), function_call)
except (AttributeError, TypeError):
pass
# Fall back to name attribute
name = getattr(function_call, "name", None)
if name:
return str(name)
return "unknown tool"

Some files were not shown because too many files have changed in this diff Show More