mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Compare commits
16
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
6b2ff3ed7a | ||
|
|
bad05a2bdc | ||
|
|
7e0767a0a0 | ||
|
|
af772997af | ||
|
|
b343625c1f | ||
|
|
9bc7b27813 | ||
|
|
6a2efeae7c | ||
|
|
6169df04cb | ||
|
|
331201294b | ||
|
|
fa9e086576 | ||
|
|
dcc218dbac | ||
|
|
6bd2cfec03 | ||
|
|
ab8ba8fc61 | ||
|
|
9cafd7e58b | ||
|
|
d5335fbeae | ||
|
|
bf4ad48cf2 |
@@ -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" />
|
||||
|
||||
+1
@@ -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
|
||||
|
||||
+19
-4
@@ -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)
|
||||
|
||||
+1
@@ -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>
|
||||
|
||||
+25
-7
@@ -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);
|
||||
}
|
||||
|
||||
+19
-6
@@ -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,
|
||||
|
||||
+10
-1
@@ -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) =>
|
||||
|
||||
+285
@@ -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 1–2: 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 1–2: 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)
|
||||
|
||||
+5
-6
@@ -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; }
|
||||
}
|
||||
+4
-4
@@ -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);
|
||||
|
||||
+46
-2
@@ -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
-3
@@ -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()
|
||||
|
||||
+1
@@ -3,6 +3,7 @@
|
||||
<PropertyGroup>
|
||||
<!-- GitHub.Copilot.SDK only supports .NET 8.0+ -->
|
||||
<TargetFrameworks>$(TargetFrameworksCore)</TargetFrameworks>
|
||||
<NoWarn>$(NoWarn);GHCP001</NoWarn>
|
||||
</PropertyGroup>
|
||||
|
||||
<ItemGroup>
|
||||
|
||||
+5
-5
@@ -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
|
||||
|
||||
+20
-21
@@ -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);
|
||||
|
||||
+1
@@ -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
|
||||
|
||||
+7
-7
@@ -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");
|
||||
|
||||
+5
-5
@@ -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");
|
||||
|
||||
+89
-6
@@ -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);
|
||||
|
||||
+574
@@ -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
|
||||
}
|
||||
+3
-3
@@ -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);
|
||||
|
||||
+310
-3
@@ -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
|
||||
}
|
||||
|
||||
+89
@@ -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."""
|
||||
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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 -----------------------------------------------------
|
||||
|
||||
@@ -7,6 +7,7 @@
|
||||
"**/demos/**",
|
||||
"**/_to_delete/**",
|
||||
"**/05-end-to-end/**",
|
||||
"**/harness/**",
|
||||
"**/agent_with_foundry_tracing.py",
|
||||
"**/azure_responses_client_with_foundry.py"
|
||||
],
|
||||
|
||||
@@ -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
Reference in New Issue
Block a user