diff --git a/docs/specs/004-python-function-calling-loop.md b/docs/specs/004-python-function-calling-loop.md index bfbe4812299..4603417996c 100644 --- a/docs/specs/004-python-function-calling-loop.md +++ b/docs/specs/004-python-function-calling-loop.md @@ -191,9 +191,12 @@ Before acting on a model function-call batch, the loop classifies every actionab call aborts the complete batch before approval state changes or execution. Otherwise approval-required and Host-owned calls are returned together in model order, while session-backed executable siblings remain deferred. An incomplete session-backed mixed approval/Host response remains pending without executing a deferred call; a stateless incomplete -response is rejected. Correlation is scoped to the active mixed batch so completed or abandoned historical Host calls -remain unchanged. `ToolApprovalMiddleware` may resolve approval requests through standing or automatic policies, but -it preserves non-approval user-input requests and does not split manual approvals away from their Host-owned siblings. +response is rejected, including when no response is supplied, message roles vary within the model output, or a newer +standalone request follows the incomplete batch. Stateless discovery selects the latest unresolved mixed batch, so a +completed batch remains inert on later turns. Correlation is scoped to the active mixed batch so completed or abandoned +historical Host calls remain unchanged. +`ToolApprovalMiddleware` may resolve approval requests through standing or automatic policies, but it preserves +non-approval user-input requests and does not split or reorder manual approvals relative to their Host-owned siblings. ### Reasoning-bound function-call groups @@ -612,7 +615,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Scenario | Required invariant | Primary regression test | |---|---|---| | Fatal call mixed with pauses | Complete-batch classification raises before approval or execution, independent of call order. | `packages/core/tests/core/test_function_invocation_logic.py::test_mixed_batch_fatal_unknown_precedes_every_pause` | -| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_active_mixed_pause_ignores_historical_host_requests` | +| Approval and Host-owned calls | Both pause types are returned in model order; a session-backed partial response remains pending across serialization; a stateless zero-response or partial response fails closed across message roles and cannot be hidden by a newer standalone request; completed mixed batches remain inert; historical Host calls do not participate; a complete response executes the exact approved arguments once. | `test_mixed_batch_returns_approval_and_host_pause_in_model_order`, `test_mixed_batch_requires_complete_responses_before_execution`, `test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution`, `test_stateless_mixed_batch_across_message_roles_requires_complete_responses`, `test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch`, `test_completed_split_stateless_mixed_batch_is_inert_on_later_turn`, `test_active_mixed_pause_ignores_historical_host_requests` | | Safe and approval-required calls in one batch | Hidden safe calls replay only with the matching visible approval. | `packages/core/tests/core/test_harness_tool_approval.py::test_mixed_batch_hides_already_approved_request_until_approval_replay` | | Restored approval state | Serialized `ToolApprovalState` restores mixed-batch behavior. | `test_mixed_batch_accepts_restored_tool_approval_state` | | Unrelated turn before approval | Hidden calls do not execute on an unrelated turn. | `test_hidden_mixed_batch_requests_do_not_replay_on_unrelated_turn` | @@ -620,7 +623,7 @@ that manually replay messages own the equivalent rule: do not resend an approval | Queued approvals | One unresolved approval is surfaced per run without premature execution. | `test_tool_approval_middleware_queues_multiple_approval_requests`, `test_tool_approval_middleware_queues_streamed_approval_requests` | | Middleware state plus hidden core state | State saves do not discard hidden mixed-batch calls. | `test_tool_approval_middleware_preserves_hidden_mixed_batch_requests` | | Auto-approval callback | Callback receives the original function call and executes the approved set once. | `test_tool_approval_middleware_auto_approval_rule_receives_function_call` | -| Approval policy with Host-owned sibling | Reordering the same calls has the same outcome; auto-approved and safe calls remain deferred until Host input arrives; manual approvals stay in the mixed batch; an approved Host tool returns to Host-owned handling before local execution. | `test_tool_approval_middleware_mixed_batch_is_order_independent`, `test_tool_approval_middleware_auto_approves_with_host_pause_and_cached_safe_call`, `test_tool_approval_middleware_keeps_manual_approvals_together_with_host_pause`, `test_tool_approval_middleware_policy_approval_reclassifies_host_tool` | +| Approval policy with Host-owned sibling | Reordering the same calls has the same outcome; auto-approved and safe calls remain deferred until Host input arrives; separately streamed Host input stays immediately visible while later manual approvals remain in model order and in the same batch; an approved Host tool returns to Host-owned handling before local execution. | `test_tool_approval_middleware_mixed_batch_is_order_independent`, `test_tool_approval_middleware_auto_approves_with_host_pause_and_cached_safe_call`, `test_tool_approval_middleware_keeps_manual_approvals_together_with_host_pause`, `test_tool_approval_middleware_preserves_split_streamed_mixed_batch_order`, `test_tool_approval_middleware_streams_standalone_user_input_immediately`, `test_tool_approval_middleware_policy_approval_reclassifies_host_tool` | | Shared call budget | Auto-approved re-entry does not reset `max_function_calls`, and every executed approval group counts even when it pauses for input. | `test_tool_approval_middleware_auto_approved_loops_share_function_call_budget`, `test_approval_resume_user_input_counts_toward_function_call_budget` | | Standing tool rule | Tool-level approval applies only to later matching tools. | `test_tool_approval_middleware_always_approve_tool_rule` | | Forged standing rule | An unbound or substituted hosted response cannot create a standing middleware approval rule for caller-selected metadata. | `test_tool_approval_middleware_drops_forged_standing_approval`, `test_tool_approval_middleware_rebinds_hosted_standing_approval` | diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs index 9e2e5824f3a..38bdd10b681 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/Models/Common/ContentToProcess.cs @@ -14,19 +14,19 @@ internal sealed class ContentToProcess /// /// Creates a new instance of ContentToProcess. /// - /// The content to send and its associated ids. + /// The content to send and its associated id. /// Metadata about the activity performed with the content. /// Metadata about the device that produced the content. /// Metadata about the application integrating with Purview. /// Metadata about the application being protected by Purview. public ContentToProcess( - List contentEntries, + ProcessContentMetadataBase contentEntry, ActivityMetadata activityMetadata, DeviceMetadata deviceMetadata, IntegratedAppMetadata integratedAppMetadata, ProtectedAppMetadata protectedAppMetadata) { - this.ContentEntries = contentEntries; + this.ContentEntries = [contentEntry]; this.ActivityMetadata = activityMetadata; this.DeviceMetadata = deviceMetadata; this.IntegratedAppMetadata = integratedAppMetadata; diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/README.md b/dotnet/src/Microsoft.Agents.AI.Purview/README.md index 39915507b25..e92d1380e13 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/README.md +++ b/dotnet/src/Microsoft.Agents.AI.Purview/README.md @@ -273,3 +273,39 @@ catch (PurviewException e) this._logger.LogError(e, "Purview middleware threw an exception.") } ``` + +## Security Considerations + +### Identity is a trusted input + +Purview evaluates DLP policy **for a specific user**. The identity this integration resolves therefore +decides *which* policy is applied, and it is resolved in this order: + +1. The user id from the configured `TokenCredential`'s token, when the credential resolves to a user. +2. The `userId` argument passed to the processor. +3. `ChatMessage.AdditionalProperties["userId"]`. +4. `ChatMessage.AuthorName`, when it is a GUID. + +Only source 1 is verified. Sources 2-4 are supplied by the hosting application, so **a host must not +populate them from data that has crossed a trust boundary**. If an end user, an upstream service or a +model response can influence `AdditionalProperties["userId"]` or `AuthorName`, that party can select a +different user's DLP policy - typically one with weaker rules - and evade enforcement. Where identity +must come from a request, derive it from a validated token on the server, never from the request body. +Prefer a user-delegated credential (source 1) whenever possible. + +`PurviewAppLocation` in `PurviewSettings` is trusted in the same way: it selects which policy locations +apply and must be configured by the host, not by the caller. + +### Fail-closed behaviour + +Policy evaluation fails closed. If no user id can be resolved, or the tenant or app location cannot be +determined, `PurviewRequestException` is thrown rather than letting content through unevaluated. Use +`IgnoreExceptions` if you deliberately want availability over enforcement - but understand that it +disables enforcement for every error, not just transient ones. + +### What is evaluated + +Every content item on a message is submitted for evaluation, not just its text: `DataContent` is sent +as Purview binary content, and `FunctionCallContent`, `FunctionResultContent` and other structured +content are serialized to text. Only `UsageContent` is skipped, because it carries token counts rather +than user data. diff --git a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs index f49f1239a89..a456f2d7891 100644 --- a/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs +++ b/dotnet/src/Microsoft.Agents.AI.Purview/ScopedContentProcessor.cs @@ -3,6 +3,7 @@ using System; using System.Collections.Generic; using System.Linq; +using System.Text.Json; using System.Threading; using System.Threading.Tasks; using Microsoft.Agents.AI.Purview.Models.Common; @@ -22,6 +23,11 @@ internal sealed class ScopedContentProcessor : IScopedContentProcessor private readonly ICacheProvider _cacheProvider; private readonly IChannelHandler _channelHandler; + /// + /// The characters that terminate the authority of a URL, being the start of the path, query or fragment. + /// + private static readonly char[] s_authorityDelimiters = ['/', '?', '#']; + /// /// Create a new instance of . /// @@ -52,12 +58,7 @@ public ScopedContentProcessor(IPurviewClient purviewClient, ICacheProvider cache foreach (DlpActionInfo policyAction in processContentResponse.PolicyActions) { // We need to process all data before blocking, so set the flag and return it outside of this loop. - if (policyAction.Action == DlpAction.BlockAccess) - { - shouldBlock = true; - } - - if (policyAction.RestrictionAction == RestrictionAction.Block) + if (IsBlockingAction(policyAction)) { shouldBlock = true; } @@ -68,6 +69,85 @@ public ScopedContentProcessor(IPurviewClient purviewClient, ICacheProvider cache return (shouldBlock, resolvedUserId); } + /// + /// Whether a policy action means the content must not be released. + /// + /// + /// is not blocking on its own: it carries a separate + /// that selects the enforcement mode, which may be + /// , or + /// as well as . Only an + /// explicit mode withholds the content. + /// + /// The policy action to inspect. + /// when the action blocks the content. + private static bool IsBlockingAction(DlpActionInfo actionInfo) + => actionInfo.Action == DlpAction.BlockAccess + || actionInfo.RestrictionAction == RestrictionAction.Block; + + /// + /// Filter policy actions down to the ones that block. + /// + /// The policy actions to filter. + /// The blocking subset of . + private static List GetBlockingActions(List actionInfos) + { + List blockingActions = []; + + foreach (DlpActionInfo actionInfo in actionInfos) + { + if (IsBlockingAction(actionInfo)) + { + blockingActions.Add(actionInfo); + } + } + + return blockingActions; + } + + /// + /// Normalize a policy location value for comparison, according to its location type. + /// + /// + /// Application ids (GUIDs) and domain names are case-insensitive, so those fold whole. URL values + /// are not. Only the scheme and the host are case-insensitive; the userinfo, path, query and + /// fragment are all case-sensitive, so folding a URL whole would let a scope for + /// contoso.com/public match a request for contoso.com/Public. Location types that are + /// not recognized fold whole, which matches more scopes rather than fewer. + /// + /// The location type segment, for example policyLocationUrl. + /// The location value to normalize. + /// The value in its comparable form. + private static string NormalizeLocationValue(string locationType, string value) + { + if (!locationType.EndsWith("url", StringComparison.OrdinalIgnoreCase)) + { + return value.ToUpperInvariant(); + } + + int schemeEnd = value.IndexOf("://", StringComparison.Ordinal); + int authorityStart = schemeEnd >= 0 ? schemeEnd + 3 : 0; + + // The authority ends at the first path, query or fragment delimiter, whichever comes first. + int authorityEnd = value.IndexOfAny(s_authorityDelimiters, authorityStart); + if (authorityEnd < 0) + { + authorityEnd = value.Length; + } + + // Credentials are case-sensitive, so only the host half of the authority folds. + string authority = value.Substring(authorityStart, authorityEnd - authorityStart); + int userInfoEnd = authority.LastIndexOf('@'); + string userInfo = userInfoEnd >= 0 ? authority.Substring(0, userInfoEnd + 1) : string.Empty; + string host = userInfoEnd >= 0 ? authority.Substring(userInfoEnd + 1) : authority; + + return string.Concat( + value.Substring(0, authorityStart).ToUpperInvariant(), + userInfo, + host.ToUpperInvariant(), + value.Substring(authorityEnd)); + } + private static bool TryGetUserIdFromPayload(IEnumerable messages, out string? userId) { userId = null; @@ -115,12 +195,9 @@ private async Task> MapMessageToPCRequestsAsync(IEnu foreach (ChatMessage message in messages) { string messageId = message.MessageId ?? Guid.NewGuid().ToString(); - ContentBase content = new PurviewTextContent(message.Text); string correlationId = (sessionId ?? Guid.NewGuid().ToString()) + "@AF"; - ProcessConversationMetadata conversationMetadata = new(content, messageId, false, $"Agent Framework Message {messageId}", correlationId) - { - SequenceNumber = DateTime.UtcNow.Ticks, - }; + long baseSequenceNumber = DateTime.UtcNow.Ticks; + List mappedContents = MapMessageContents(message); ActivityMetadata activityMetadata = new(activity); PolicyLocation policyLocation; @@ -158,20 +235,127 @@ private async Task> MapMessageToPCRequestsAsync(IEnu OperatingSystemVersion = "Unknown" } }; - ContentToProcess contentToProcess = new([conversationMetadata], activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata); if (string.IsNullOrEmpty(resolvedUserId)) { throw new PurviewRequestException("No user id provided or inferred for Purview request. Please provide an Entra user id in each message, pass a user id to the processor, or configure the TokenCredential to authenticate to an Entra user."); } - ProcessContentRequest pcRequest = new(contentToProcess, resolvedUserId, tenantId); - pcRequests.Add(pcRequest); + for (int entryIndex = 0; entryIndex < mappedContents.Count; entryIndex++) + { + string identifier = entryIndex == 0 ? messageId : $"{messageId}-{entryIndex}"; + ProcessConversationMetadata contentEntry = new(mappedContents[entryIndex], identifier, false, $"Agent Framework Message {messageId}", correlationId) + { + SequenceNumber = baseSequenceNumber + entryIndex, + }; + ContentToProcess contentToProcess = new(contentEntry, activityMetadata, deviceMetadata, integratedAppMetadata, protectedAppMetadata); + ProcessContentRequest pcRequest = new(contentToProcess, resolvedUserId, tenantId); + pcRequests.Add(pcRequest); + } } return pcRequests; } + /// + /// Map every content item of a message onto the Purview content type that fits it. + /// + /// The message whose contents should be evaluated. + /// One content item per evaluable . + /// + /// and content with no data are skipped because they carry no user + /// data. Everything else is mapped to a real content item: submitting a message with part of + /// its payload unevaluated is a policy bypass. + /// + private static List MapMessageContents(ChatMessage message) + { + List mapped = []; + + foreach (AIContent content in message.Contents) + { + ContentBase? purviewContent = MapContent(content); + if (purviewContent != null) + { + mapped.Add(purviewContent); + } + } + + return mapped; + } + + /// + /// Map a single onto a Purview content item. + /// + /// The content item to map. + /// The Purview content item, or when the content carries no user data. + private static ContentBase? MapContent(AIContent content) + { + switch (content) + { + case UsageContent: + // Telemetry only; there is nothing for DLP to classify. + return null; + + case TextContent { Text: null or "" }: + case TextReasoningContent { Text: null or "" }: + case DataContent { Data.IsEmpty: true }: + // Empty data is skipped because the Graph APIs do not support it. A request with + // empty data cannot violate content policies because it contains no content. + return null; + + case TextContent textContent: + return new PurviewTextContent(textContent.Text); + + case TextReasoningContent reasoningContent: + return new PurviewTextContent(reasoningContent.Text); + + case DataContent dataContent: + return new PurviewBinaryContent(dataContent.Data.ToArray()); + + case UriContent uriContent: + return new PurviewTextContent($"{uriContent.MediaType} {uriContent.Uri}"); + + case FunctionCallContent functionCallContent: + return new PurviewTextContent(SerializeForEvaluation(functionCallContent)); + + case FunctionResultContent functionResultContent: + return new PurviewTextContent(SerializeForEvaluation(functionResultContent)); + + case ErrorContent errorContent: + return new PurviewTextContent(errorContent.Message ?? string.Empty); + + default: + // Catch-all for every remaining content type, including ones added after this code + // was written. Serializing is always safe; skipping would be a policy bypass. + return new PurviewTextContent(SerializeForEvaluation(content)); + } + } + + /// + /// Render a content item as text so Purview can classify everything it carries. + /// + /// The content item to render. + /// The text representation of . + /// The content could not be rendered for evaluation. + /// + /// The type info is resolved from the concrete runtime type rather than . + /// Polymorphic serialization through the base type throws for any subclass the abstractions do not + /// declare, which would push third-party and future content types onto the failure path. + /// + private static string SerializeForEvaluation(AIContent content) + { + try + { + return JsonSerializer.Serialize(content, AIJsonUtilities.DefaultOptions.GetTypeInfo(content.GetType())); + } + catch (Exception ex) when (ex is NotSupportedException or InvalidOperationException or JsonException) + { + // Fail closed. Falling back to ToString() would submit a bare type name in place of the + // payload, so the content would clear policy without ever having been evaluated. + throw new PurviewRequestException($"Unable to serialize content of type '{content.GetType()}' for Purview policy evaluation. Content cannot be evaluated and will not be released.", ex); + } + } + /// /// Orchestrates process content and protection scopes calls. /// @@ -228,12 +412,18 @@ private async Task ProcessWithCachedScopesAsync( if (shouldProcess) { - pcRequest.ProcessInline = executionMode == ExecutionMode.EvaluateInline; + pcRequest.ProcessInline = executionMode != ExecutionMode.EvaluateOffline; if (executionMode == ExecutionMode.EvaluateOffline) { this._channelHandler.QueueJob(new ProcessContentJob(pcRequest)); - return new ProcessContentResponse(); + + // Offline evaluation is asynchronous, but an explicit block action already known + // from the cached scopes must still be enforced rather than discarded. + List blockingActions = GetBlockingActions(dlpActions); + return blockingActions.Count > 0 + ? new ProcessContentResponse { PolicyActions = blockingActions } + : new ProcessContentResponse(); } return await this.CallProcessContentAsync(pcRequest, cacheKey, dlpActions, cancellationToken).ConfigureAwait(false); @@ -327,7 +517,8 @@ internal static (bool shouldProcess, List dlpActions, ExecutionMo foreach (var location in scope.Locations ?? Array.Empty()) { - if (location.DataType.EndsWith(locationType, StringComparison.OrdinalIgnoreCase) && location.Value.Equals(locationValue, StringComparison.OrdinalIgnoreCase)) + if (location.DataType.EndsWith(locationType, StringComparison.OrdinalIgnoreCase) + && NormalizeLocationValue(locationType, location.Value).Equals(NormalizeLocationValue(locationType, locationValue), StringComparison.Ordinal)) { locationMatch = true; break; @@ -338,7 +529,9 @@ internal static (bool shouldProcess, List dlpActions, ExecutionMo { shouldProcess = true; - if (scope.ExecutionMode == ExecutionMode.EvaluateInline) + // Only an explicitly offline scope may skip inline evaluation. ExecutionMode is an + // evolvable enum, so any unrecognised value is upgraded to inline (fail closed). + if (scope.ExecutionMode != ExecutionMode.EvaluateOffline) { executionMode = ExecutionMode.EvaluateInline; } diff --git a/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs b/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs index 136d4d6833a..06e0e7853da 100644 --- a/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs +++ b/dotnet/src/Microsoft.Agents.AI/Skills/AgentSkillFrontmatter.cs @@ -112,6 +112,10 @@ public string? Compatibility /// /// Gets or sets the arbitrary key-value metadata for this skill. /// + /// + /// Keys are compared case-insensitively. When parsed from a SKILL.md file, duplicate + /// entries retain the first value and key spelling and produce warnings without rejecting the skill. + /// public AdditionalPropertiesDictionary? Metadata { get; set; } /// diff --git a/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs b/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs index b22ea74c93e..5484348c59e 100644 --- a/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs +++ b/dotnet/src/Microsoft.Agents.AI/Skills/File/AgentFileSkillsSource.cs @@ -22,7 +22,10 @@ namespace Microsoft.Agents.AI; /// /// Searches directories recursively (up to 2 levels deep) for SKILL.md files. /// Symbolic links and reparse points below configured roots are not followed during skill discovery. -/// Each file is validated for YAML frontmatter. Resource and script files are discovered by scanning the skill +/// Recognized top-level frontmatter fields must use lowercase names and must not be repeated. +/// Within the optional metadata mapping, keys are compared case-insensitively. The first +/// value is retained for duplicate keys, and subsequent entries produce warnings without rejecting the skill. +/// Resource and script files are discovered by scanning the skill /// directory for files with matching extensions. Invalid resources are skipped with logged warnings. /// Resource and script paths are checked against path traversal and symlink escape attacks. /// Discovered files are revalidated against their trusted skill directory immediately before use. @@ -36,19 +39,31 @@ public sealed partial class AgentFileSkillsSource : AgentSkillsSource private static readonly string[] s_defaultScriptExtensions = [".py", ".js", ".sh", ".ps1", ".cs", ".csx"]; private static readonly string[] s_defaultResourceExtensions = [".md", ".json", ".yaml", ".yml", ".csv", ".xml", ".txt"]; + // Case-insensitive lookup identifies known fields; values preserve the required spelling for validation. + private static readonly Dictionary s_frontmatterFieldNames = new(StringComparer.OrdinalIgnoreCase) + { + ["name"] = "name", + ["description"] = "description", + ["license"] = "license", + ["compatibility"] = "compatibility", + ["metadata"] = "metadata", + ["allowed-tools"] = "allowed-tools", + }; + // Matches YAML frontmatter delimited by "---" lines. Group 1 = content between delimiters. // Multiline makes ^/$ match line boundaries; Singleline makes . match newlines across the block. // The \uFEFF? prefix allows an optional UTF-8 BOM that some editors prepend. private static readonly Regex s_frontmatterRegex = new(@"\A\uFEFF?^---\s*$(.+?)^---\s*$", RegexOptions.Multiline | RegexOptions.Singleline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); - // Matches top-level YAML "key: value" lines. Group 1 = key (supports hyphens for keys like allowed-tools), - // Group 2 = quoted value, Group 3 = unquoted value. - // Accepts single or double quotes; the lazy quantifier trims trailing whitespace on unquoted values. - private static readonly Regex s_yamlKeyValueRegex = new(@"^([\w-]+)\s*:\s*(?:[""'](.+?)[""']|(.+?))\s*$", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); + // Matches top-level YAML "key: value" lines. Group 1 = key (including quotes), + // Group 2 = quoted value, Group 3 = unquoted value (possibly empty for keys such as "metadata:"). + // A value may start on a later indented line, but cannot consume another top-level field. + // Retain trailing whitespace matching so block-scalar parsing handles leading blank lines as before. + private static readonly Regex s_yamlKeyValueRegex = new(@"^([\w-]+|""[\w-]+""|'[\w-]+')[ \t]*:[ \t]*(?:\r?\n(?:[ \t]*\r?\n)*[ \t]+)?(?:[""']([^\r\n]+?)[""']|([^\r\n]*?))\s*$", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); - // Matches a "metadata:" line followed by indented sub-key/value pairs. + // Matches a metadata block, allowing quotes around the root key. // Group 1 captures the entire indented block beneath the metadata key. - private static readonly Regex s_yamlMetadataBlockRegex = new(@"^metadata\s*:\s*$\n((?:[ \t]+\S.*\n?)+)", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); + private static readonly Regex s_yamlMetadataBlockRegex = new(@"^(?:metadata|""metadata""|'metadata')\s*:\s*$\n((?:[ \t]+\S.*\n?)+)", RegexOptions.Multiline | RegexOptions.Compiled, TimeSpan.FromSeconds(5)); // Matches indented YAML "key: value" lines within a metadata block. // Group 1 = key (supports hyphens), Group 2 = quoted value, Group 3 = unquoted value. @@ -235,30 +250,59 @@ private bool TryParseFrontmatter(string content, string skillFilePath, [NotNullW string? compatibility = null; string? allowedTools = null; + // Recognized fields use exact lowercase names, with optional quotes. Reject casing variants and + // duplicates rather than silently changing which value the skill exposes. + var seenFields = new HashSet(StringComparer.Ordinal); foreach (Match kvMatch in s_yamlKeyValueRegex.Matches(yamlContent)) { string key = kvMatch.Groups[1].Value; + key = key[0] is '"' or '\'' ? key.Substring(1, key.Length - 2) : key; + + // Unknown fields are intentionally excluded from this validation for forward compatibility. + if (!s_frontmatterFieldNames.TryGetValue(key, out string? canonicalKey)) + { + continue; + } + + if (!string.Equals(key, canonicalKey, StringComparison.Ordinal)) + { + LogIncorrectlyCasedFrontmatterField(this._logger, skillFilePath, key, canonicalKey); + return false; + } + + if (!seenFields.Add(key)) + { + LogDuplicateFrontmatterField(this._logger, skillFilePath, key); + return false; + } + + // Empty declarations participate in key validation, but leave optional scalar fields unset. + if (!kvMatch.Groups[2].Success && kvMatch.Groups[3].Length == 0) + { + continue; + } + string value = kvMatch.Groups[2].Success ? kvMatch.Groups[2].Value : ParseYamlScalarValue(yamlContent, kvMatch); - if (string.Equals(key, "name", StringComparison.OrdinalIgnoreCase)) + if (string.Equals(key, "name", StringComparison.Ordinal)) { name = value; } - else if (string.Equals(key, "description", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "description", StringComparison.Ordinal)) { description = value; } - else if (string.Equals(key, "license", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "license", StringComparison.Ordinal)) { license = value; } - else if (string.Equals(key, "compatibility", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "compatibility", StringComparison.Ordinal)) { compatibility = value; } - else if (string.Equals(key, "allowed-tools", StringComparison.OrdinalIgnoreCase)) + else if (string.Equals(key, "allowed-tools", StringComparison.Ordinal)) { allowedTools = value; } @@ -272,7 +316,14 @@ private bool TryParseFrontmatter(string content, string skillFilePath, [NotNullW metadata = []; foreach (Match kvMatch in s_yamlIndentedKeyValueRegex.Matches(metadataMatch.Groups[1].Value)) { - metadata[kvMatch.Groups[1].Value] = kvMatch.Groups[2].Success ? kvMatch.Groups[2].Value : kvMatch.Groups[3].Value; + string key = kvMatch.Groups[1].Value; + string value = kvMatch.Groups[2].Success ? kvMatch.Groups[2].Value : kvMatch.Groups[3].Value; + + // Keep the first value and key spelling using the dictionary's case-insensitive comparison. + if (!metadata.TryAdd(key, value)) + { + LogDuplicateMetadataKey(this._logger, skillFilePath, key); + } } } @@ -722,6 +773,15 @@ private static void ValidateExtensions(IEnumerable? extensions) [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}' has an invalid '{FieldName}' value: {Reason}")] private static partial void LogInvalidFieldValue(ILogger logger, string skillFilePath, string fieldName, string reason); + [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}' contains duplicate frontmatter field '{FieldName}'")] + private static partial void LogDuplicateFrontmatterField(ILogger logger, string skillFilePath, string fieldName); + + [LoggerMessage(LogLevel.Warning, "SKILL.md at '{SkillFilePath}' contains duplicate metadata key '{Key}'; keeping the first value")] + private static partial void LogDuplicateMetadataKey(ILogger logger, string skillFilePath, string key); + + [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}' uses incorrectly cased frontmatter field '{FieldName}'; expected '{ExpectedFieldName}'")] + private static partial void LogIncorrectlyCasedFrontmatterField(ILogger logger, string skillFilePath, string fieldName, string expectedFieldName); + [LoggerMessage(LogLevel.Error, "SKILL.md at '{SkillFilePath}': skill name '{SkillName}' does not match parent directory name '{DirectoryName}'")] private static partial void LogNameDirectoryMismatch(ILogger logger, string skillFilePath, string skillName, string directoryName); diff --git a/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs b/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs index 418ecd9c19c..f9754c34960 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Mcp.UnitTests/Skills/AgentMcpSkillsSourceArchiveTests.cs @@ -62,6 +62,186 @@ Content B. private const int ManyFileArchiveFileCount = 60; + [Theory] + [InlineData("description: First\ndescription: >-\n Second", "\n")] + [InlineData("description: |-\n First\nDescription: Second", "\n")] + [InlineData("description: Valid\nmetadata:\n author: First\nmetadata:\n author: Second", "\n")] + [InlineData("description: Valid\nAllowed-Tools: read", "\n")] + [InlineData("description: Valid\nallowed-tools: read\n\"allowed-tools\": other", "\n")] + [InlineData("description: Valid\nallowed-tools: read\n\"allowed-tools\": other", "\r\n")] + [InlineData("description: Valid\n'allowed-tools': other\nallowed-tools: read", "\n")] + [InlineData("description: Valid\n'allowed-tools': other\nallowed-tools: read", "\r\n")] + [InlineData("description: Valid\n\"description\": Other text", "\n")] + [InlineData("'description': Other text\ndescription: Valid", "\r\n")] + [InlineData("description: Valid\n\"Allowed-Tools\": other", "\n")] + [InlineData("description: Valid\nmetadata:\n author: First\n\"metadata\":\n author: Second", "\r\n")] + [InlineData("description: Valid\nlicense: MIT\n'license':", "\n")] + [InlineData("description: Valid\ncompatibility: Any runtime\n\"compatibility\": Other runtime", "\r\n")] + [InlineData("description: Valid\n\"name\": archived-skill", "\n")] + public async Task GetSkillsAsync_AmbiguousArchiveFrontmatter_SkipsSkillAsync(string fields, string newline) + { + // Arrange + string content = $"---\nname: archived-skill\n{fields}\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + + [Theory] + [InlineData("\n")] + [InlineData("\r\n")] + public async Task GetSkillsAsync_ArchiveValueOnIndentedNextLine_PreservedAsync(string newline) + { + // Arrange + const string Content = "---\nname: archived-skill\ndescription:\n 'Read files'\nlicense: MIT\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", Content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Equal("MIT", frontmatter.License); + } + + [Theory] + [InlineData("license", "\n")] + [InlineData("license", "\r\n")] + [InlineData("compatibility", "\n")] + [InlineData("compatibility", "\r\n")] + [InlineData("allowed-tools", "\n")] + [InlineData("allowed-tools", "\r\n")] + public async Task GetSkillsAsync_ArchiveEmptyOptionalScalar_RemainsNullAsync(string field, string newline) + { + // Arrange + string content = $"---\nname: archived-skill\ndescription: Read files\n{field}: \t\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Null(frontmatter.License); + Assert.Null(frontmatter.Compatibility); + Assert.Null(frontmatter.AllowedTools); + } + + [Theory] + [InlineData("author", "\n", "")] + [InlineData("author", "\r\n", "")] + [InlineData("Author", "\n", "")] + [InlineData("Author", "\r\n", "")] + [InlineData("author", "\n", "'")] + [InlineData("author", "\r\n", "'")] + [InlineData("Author", "\n", "'")] + [InlineData("Author", "\r\n", "'")] + [InlineData("author", "\n", "\"")] + [InlineData("author", "\r\n", "\"")] + [InlineData("Author", "\n", "\"")] + [InlineData("Author", "\r\n", "\"")] + public async Task GetSkillsAsync_ArchiveDuplicateMetadata_KeepsFirstValueAndLogsWarningAsync( + string secondKey, string newline, string quote) + { + // Arrange + string content = + $"---\n{quote}name{quote}: archived-skill\n{quote}description{quote}: >-\n Read\n files\n" + + $"{quote}compatibility{quote}:\n Any runtime\n{quote}allowed-tools{quote}: 'read'\n{quote}metadata{quote}:\n" + + $" author:\n 'First'\n {secondKey}: Second\n author: Third\n version: '1.0'\n" + + $"{quote}license{quote}: |-\n MIT\n License\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentMcpSkillsSource( + client, new() { ArchiveSkillsDirectory = this._extractionRoot }, loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var skill = Assert.Single(skills); + Assert.NotNull(skill.Frontmatter.Metadata); + Assert.Equal(2, skill.Frontmatter.Metadata.Count); + Assert.Equal("First", skill.Frontmatter.Metadata["author"]); + Assert.Equal("First", skill.Frontmatter.Metadata["Author"]); + Assert.Equal("1.0", skill.Frontmatter.Metadata["version"]); + Assert.Equal("Read files", skill.Frontmatter.Description); + Assert.Equal("MIT\nLicense", skill.Frontmatter.License); + Assert.Equal("Any runtime", skill.Frontmatter.Compatibility); + Assert.Equal("read", skill.Frontmatter.AllowedTools); + Assert.Contains("Body.", await skill.GetContentAsync()); + logger.Verify( + l => l.Log( + LogLevel.Warning, + It.IsAny(), + It.Is((state, _) => + state.ToString()!.Contains("duplicate metadata key", StringComparison.Ordinal) && + state.ToString()!.Contains("keeping the first value", StringComparison.Ordinal)), + It.IsAny(), + It.IsAny>()), + Times.Exactly(2)); + } + + [Theory] + [InlineData("First", "First", "\n")] + [InlineData("First", "First", "\r\n")] + [InlineData("\n 'First'", "First", "\n")] + [InlineData("\n 'First'", "First", "\r\n")] + [InlineData("|-\n First\n Second", "|-", "\n")] + [InlineData("|-\n First\n Second", "|-", "\r\n")] + [InlineData(">-\n First\n Second", ">-", "\n")] + [InlineData(">-\n First\n Second", ">-", "\r\n")] + public async Task GetSkillsAsync_ArchiveExistingMetadataScalarFormats_PreservedAsync(string value, string expected, string newline) + { + // Arrange + string content = "---\nname: archived-skill\ndescription: Read files\n" + + $"metadata:\n author: {value}\n version: '1.0'\n---\nBody."; + byte[] archive = BuildZip(("SKILL.md", content.Replace("\n", newline))); + await using var server = CreateArchiveServer( + ArchiveIndex("archived-skill", "skill://archives/archived-skill.zip"), + new Dictionary { ["archived-skill"] = archive }); + await using var client = await server.CreateClientAsync(); + var source = new AgentMcpSkillsSource(client, new() { ArchiveSkillsDirectory = this._extractionRoot }); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert - preserve the existing metadata scalar representation. + var skill = Assert.Single(skills); + Assert.NotNull(skill.Frontmatter.Metadata); + Assert.Equal(2, skill.Frontmatter.Metadata.Count); + Assert.Equal(expected, skill.Frontmatter.Metadata["author"]); + Assert.Equal("1.0", skill.Frontmatter.Metadata["version"]); + Assert.Contains("Body.", await skill.GetContentAsync()); + } + /// /// Malformed, unsupported, and mismatched archive digests that must be rejected. /// diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs index c3415d42af4..fa9c19878ca 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/PurviewClientTests.cs @@ -519,7 +519,7 @@ private static ContentToProcess CreateValidContentToProcess() }; return new ContentToProcess( - [metadata], + metadata, activityMetadata, deviceMetadata, integratedAppMetadata, diff --git a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs index b59cc7a3929..7afad015ca1 100644 --- a/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.Purview.UnitTests/ScopedContentProcessorTests.cs @@ -296,6 +296,175 @@ public void CheckApplicableScopes_MatchesAnyLocationInScope() Assert.Equal(ExecutionMode.EvaluateInline, executionMode); } + [Fact] + public void CheckApplicableScopes_MatchesUrlLocationHostCaseInsensitively() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "HTTPS://Contoso.com/sites/marketing"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com/sites/marketing")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, _, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.True(shouldProcess); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationPathAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://contoso.com/sites/Marketing"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com/sites/marketing")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationQueryAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://contoso.com?label=Secret"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com?label=secret")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationFragmentAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://contoso.com#Section"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://contoso.com#section")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_TreatsUrlLocationUserInfoAsCaseSensitive() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationUrl", "https://alice:SecretPass@contoso.com/docs"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = [new("#microsoft.graph.policyLocationUrl", "https://alice:secretpass@contoso.com/docs")], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, List dlpActions, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.False(shouldProcess); + Assert.Empty(dlpActions); + } + + [Fact] + public void CheckApplicableScopes_MatchesApplicationLocationCaseInsensitively() + { + // Arrange + ProcessContentRequest pcRequest = CreateProcessContentRequest(); + pcRequest.ContentToProcess.ProtectedAppMetadata.ApplicationLocation = + new("microsoft.graph.policyLocationApplication", "A1B2C3D4-E5F6-4A5B-8C9D-0E1F2A3B4C5D"); + ProtectionScopesResponse psResponse = new() + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("#microsoft.graph.policyLocationApplication", "a1b2c3d4-e5f6-4a5b-8c9d-0e1f2a3b4c5d") + ], + ExecutionMode = ExecutionMode.EvaluateInline + } + ] + }; + + // Act + (bool shouldProcess, _, _) = ScopedContentProcessor.CheckApplicableScopes(pcRequest, psResponse); + + // Assert + Assert.True(shouldProcess); + } + [Fact] public async Task ProcessMessagesAsync_UsesCachedProtectionScopes_WhenAvailableAsync() { @@ -993,10 +1162,287 @@ public async Task BackgroundJobRunner_ScopeRetrievalPaymentRequired_CachesForSub It.IsAny()), Times.Once); } + /// + /// Verifies every content item is submitted for evaluation in a separate request, not just message.Text. + /// Images, binary payloads and structured tool results are empty when flattened to text, which + /// would have them reach the model having only ever been classified as an empty string. + /// + [Fact] + public async Task ProcessMessagesAsync_WithMultipleContentItems_SubmitsEachInSeparateRequestAsync() + { + // Arrange + byte[] secret = [0x01, 0x02, 0x03, 0x04]; + var messages = new List + { + new(ChatRole.User, + [ + new DataContent(secret, "application/octet-stream"), + new FunctionCallContent("call-1", "exfiltrate", new Dictionary { ["ssn"] = "123-45-6789" }) + ]) + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); + + List capturedRequests = []; + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .Callback((request, _) => capturedRequests.Add(request)) + .ReturnsAsync(new ProcessContentResponse()); + + // Act + await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.Equal(2, capturedRequests.Count); + ProcessContentMetadataBase binaryEntry = Assert.Single(capturedRequests[0].ContentToProcess.ContentEntries); + ProcessContentMetadataBase functionCallEntry = Assert.Single(capturedRequests[1].ContentToProcess.ContentEntries); + + PurviewBinaryContent binaryContent = Assert.IsType(binaryEntry.Content); + Assert.Equal(secret, binaryContent.Data); + + PurviewTextContent functionCallContent = Assert.IsType(functionCallEntry.Content); + Assert.Contains("123-45-6789", functionCallContent.Data, StringComparison.Ordinal); + + // Content entries must be individually addressable, not collapsed onto one identifier. + Assert.NotEqual(binaryEntry.Identifier, functionCallEntry.Identifier); + } + + [Fact] + public async Task ProcessMessagesAsync_WithEmptyTextAndData_SkipsEmptyContentAsync() + { + // Arrange + byte[] emptyData = []; + List messages = + [ + new(ChatRole.User, + [ + new TextContent(string.Empty), + new DataContent(emptyData, "application/octet-stream") + ]) + ]; + PurviewSettings settings = CreateValidPurviewSettings(); + TokenInfo tokenInfo = new() { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + // Act + await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Never); + } + + /// + /// Verifies a block verdict known from an offline cached scope is still enforced. The offline + /// branch reports asynchronously, but discarding the scope's own policy actions would let a + /// known restrictAccess verdict go unenforced. + /// + [Fact] + public async Task ProcessMessagesAsync_WithOfflineScopeCarryingBlockAction_ReturnsShouldBlockTrueAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + var psResponse = new ProtectionScopesResponse + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ], + ExecutionMode = ExecutionMode.EvaluateOffline, + PolicyActions = + [ + new() { Action = DlpAction.RestrictAccess, RestrictionAction = RestrictionAction.Block } + ] + } + ] + }; + + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(psResponse); + + // Act + var result = await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.True(result.shouldBlock); + + // The offline report is still queued; enforcement is additional, not a replacement. + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + } + + /// + /// Verifies a restrictAccess scope whose restriction mode does not block is not enforced as + /// a block. The action carries a separate that may be + /// , or + /// , none of which withhold the content. + /// + [Fact] + public async Task ProcessMessagesAsync_WithOfflineScopeCarryingAuditRestriction_ReturnsShouldBlockFalseAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + + var psResponse = new ProtectionScopesResponse + { + Scopes = + [ + new() + { + Activities = ProtectionScopeActivities.UploadText, + Locations = + [ + new("microsoft.graph.policyLocationApplication", "app-123") + ], + ExecutionMode = ExecutionMode.EvaluateOffline, + PolicyActions = + [ + new() { Action = DlpAction.RestrictAccess, RestrictionAction = RestrictionAction.Audit } + ] + } + ] + }; + + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(psResponse); + + // Act + var result = await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.False(result.shouldBlock); + + // The offline report is still queued so the audit action is recorded. + this._mockChannelHandler.Verify(x => x.QueueJob(It.IsAny()), Times.Once); + } + + /// + /// Verifies an unrecognised is evaluated inline. The enum is + /// evolvable, so an unknown member must not be treated as permission to skip enforcement. + /// + [Fact] + public async Task ProcessMessagesAsync_WithUnknownExecutionMode_EvaluatesInlineAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, "Test message") + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse((ExecutionMode)9999)); + + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(new ProcessContentResponse + { + PolicyActions = [new() { Action = DlpAction.BlockAccess }] + }); + + // Act + var result = await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.True(result.shouldBlock); + this._mockPurviewClient.Verify(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny()), Times.Once); + } + + /// + /// Verifies an subclass the abstractions do not declare still has its + /// payload evaluated. Serializing such a type through the polymorphic base type throws, and + /// substituting its type name would submit no payload at all for classification. + /// + [Fact] + public async Task ProcessMessagesAsync_WithUnknownContentSubclass_SubmitsPayloadForEvaluationAsync() + { + // Arrange + var messages = new List + { + new(ChatRole.User, [new UnknownTestContent { Secret = "ssn 123-45-6789" }]) + }; + var settings = CreateValidPurviewSettings(); + var tokenInfo = new TokenInfo { TenantId = "tenant-123", UserId = "user-123", ClientId = "client-123" }; + + this._mockPurviewClient.Setup(x => x.GetUserInfoFromTokenAsync(It.IsAny(), null)) + .ReturnsAsync(tokenInfo); + this._mockCacheProvider.Setup(x => x.GetAsync( + It.IsAny(), It.IsAny())) + .ReturnsAsync(CreateApplicableProtectionScopesResponse()); + + ProcessContentRequest? capturedRequest = null; + this._mockPurviewClient.Setup(x => x.ProcessContentAsync( + It.IsAny(), It.IsAny())) + .Callback((request, _) => capturedRequest = request) + .ReturnsAsync(new ProcessContentResponse()); + + // Act + await this._processor.ProcessMessagesAsync( + messages, "session-123", Activity.UploadText, settings, "user-123", CancellationToken.None); + + // Assert + Assert.NotNull(capturedRequest); + ProcessContentMetadataBase entry = Assert.Single(capturedRequest.ContentToProcess.ContentEntries); + PurviewTextContent content = Assert.IsType(entry.Content); + + Assert.Contains("123-45-6789", content.Data, StringComparison.Ordinal); + Assert.NotEqual(nameof(UnknownTestContent), content.Data); + } + #endregion #region Helper Methods + /// + /// Stands in for a third-party or future subclass that the abstractions do + /// not declare a polymorphic discriminator for. + /// + private sealed class UnknownTestContent : AIContent + { + public string Secret { get; set; } = string.Empty; + } + private static ProtectionScopesRequest CreateProtectionScopesRequest() { return new ProtectionScopesRequest("user-123", "tenant-123") @@ -1054,7 +1500,7 @@ private static ProcessContentRequest CreateProcessContentRequest() Version = "1.0" }; ContentToProcess contentToProcess = new( - [metadata], + metadata, activityMetadata, deviceMetadata, integratedAppMetadata, diff --git a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs index 6db43554f84..1599023e1b7 100644 --- a/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs +++ b/dotnet/tests/Microsoft.Agents.AI.UnitTests/AgentSkills/FileAgentSkillLoaderTests.cs @@ -12,6 +12,8 @@ using System.Security.AccessControl; using System.Security.Principal; using System.Threading.Tasks; +using Microsoft.Extensions.Logging; +using Moq; namespace Microsoft.Agents.AI.UnitTests.AgentSkills; @@ -41,6 +43,446 @@ public void Dispose() } } + /// + /// Existing scalar representations with metadata before or after them, with and without duplicate metadata. + /// + public static TheoryData ExistingScalarFormats + { + get + { + var data = new TheoryData(); + // Pin the existing lightweight parser's output, not full YAML conformance. + (string Value, string Expected)[] cases = + [ + ("Read files", "Read files"), + ("'Read files'", "Read files"), + ("\"Read files\"", "Read files"), + ("\" Read files \"", " Read files "), + ("\"Use #tags: safely\"", "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", "Use \"quotes\" safely"), + ("\"| not a block\"", "| not a block"), + ("\n Read files", "Read files"), + ("\n\n 'Read files'", "Read files"), + ("\n >-\n Read\n files", "Read files"), + ("|\n Read\n files", "Read\nfiles"), + ("|-\n Read\n files", "Read\nfiles"), + ("|+\n Read\n files", "Read\nfiles\n"), + (">\n Read\n files", "Read files"), + (">-\n Read\n files", "Read files"), + (">+\n Read\n files", "Read files\n"), + ("|-\n\n Read\n files", "Read\nfiles"), + ("|-\n Read\n\n files", "Read\n\nfiles"), + ("|-\n Read\n indented\n files", "Read\n indented\nfiles"), + ("|-\n description: text\n Description: text", "description: text\nDescription: text"), + ("|-\n \"description\": text\n 'allowed-tools': text", "\"description\": text\n'allowed-tools': text"), + ]; + foreach (var (value, expected) in cases) + { + foreach (string newline in new[] { "\n", "\r\n" }) + { + foreach (bool metadataFirst in new[] { false, true }) + { + foreach (string quote in new[] { "", "'", "\"" }) + { + data.Add(value, expected, newline, metadataFirst, false, quote); + data.Add(value, expected, newline, metadataFirst, true, quote); + } + } + } + } + + return data; + } + } + + /// + /// Duplicate and incorrectly cased root properties with quoted names. + /// + public static TheoryData ConflictingQuotedFields + { + get + { + var data = new TheoryData(); + (string Key, string Value)[] fields = + [ + ("name", "test-skill"), + ("description", "Read files"), + ("license", "MIT"), + ("compatibility", "Any runtime"), + ("metadata", "\n author: First"), + ("allowed-tools", "read"), + ]; + foreach (var (key, value) in fields) + { + string otherFields = string.Join("\n", fields.Where(f => f.Key != key).Select(f => $"{f.Key}: {f.Value}")); + foreach (char quote in new[] { '\'', '"' }) + { + string quotedKey = $"{quote}{key}{quote}"; + string quotedField = $"{quotedKey}: {value}"; + string bareField = $"{key}: {value}"; + char otherQuote = quote == '\'' ? '"' : '\''; + foreach (string newline in new[] { "\n", "\r\n" }) + { + string incorrectlyCasedField = $"{quote}{key.ToUpperInvariant()}{quote}: {value}"; + string content = $"---\n{otherFields}\n{incorrectlyCasedField}\n---\nBody."; + string diagnostic = $"incorrectly cased frontmatter field '{key.ToUpperInvariant()}'; expected '{key}'"; + data.Add(content.Replace("\n", newline), diagnostic); + foreach (string declaration in new[] + { + $"{quotedField}\n{bareField}", + $"{bareField}\n{quotedField}", + $"{quotedField}\n{quotedField}", + $"{quotedField}\n{otherQuote}{key}{otherQuote}: {value}", + $"{quotedKey}:\n{bareField}", + $"{bareField}\n{quotedKey}:", + }) + { + content = $"---\n{otherFields}\n{declaration}\n---\nBody."; + data.Add(content.Replace("\n", newline), $"duplicate frontmatter field '{key}'"); + } + } + } + } + + return data; + } + } + + /// + /// Existing metadata scalar outputs, including the lightweight parser's incomplete YAML normalization. + /// + public static TheoryData ExistingMetadataScalarFormats + { + get + { + var data = new TheoryData(); + (string Value, string Expected)[] cases = + [ + ("First", "First"), + ("'First'", "First"), + ("\"First\"", "First"), + ("\" First \"", " First "), + ("\"Use #tags: safely\"", "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", "Use \"quotes\" safely"), + ("''", "''"), + ("\"\"", "\"\""), + ("\"\\n\"", "\\n"), + ("'It''s fine'", "It''s fine"), + ("\n First", "First"), + ("\n 'First'", "First"), + ("\n \"First\"", "First"), + ("First\n Second", "First"), + ("\"First\n Second\"", "\"First"), + ("'First\n Second'", "'First"), + ("|\n First\n Second", "|"), + ("|-\n First\n Second", "|-"), + ("|+\n First\n Second", "|+"), + (">\n First\n Second", ">"), + (">-\n First\n Second", ">-"), + (">+\n First\n Second", ">+"), + ("\n |-\n First\n Second", "|-"), + ]; + foreach (var (value, expected) in cases) + { + data.Add(value, expected, "\n"); + data.Add(value, expected, "\r\n"); + } + + return data; + } + } + + /// + /// Duplicate and incorrectly cased fields before and after inline or block values. + /// + public static TheoryData ConflictingScalarFormats + { + get + { + var data = new TheoryData(); + string[] firstValues = ["First", "'First'", "\"First\"", "|-\n First", ">-\n First"]; + string[] secondValues = ["Second", "'Second'", "\"Second\"", "|-\n Second", ">-\n Second"]; + string[] secondKeys = ["description", "Description", "DESCRIPTION"]; + foreach (string firstValue in firstValues) + { + foreach (string secondValue in secondValues) + { + foreach (string secondKey in secondKeys) + { + data.Add(firstValue, secondValue, secondKey); + } + } + } + + return data; + } + } + + /// + /// Duplicate metadata values and case-insensitive key collisions, with LF and CRLF. + /// + public static TheoryData DuplicateMetadataEntries + { + get + { + var data = new TheoryData(); + (string Key, string DuplicateKey, string Value, string DuplicateValue, string Expected)[] cases = + [ + ("author", "author", "First", "Second", "First"), + ("Author", "Author", "'First'", "\"Second\"", "First"), + ("author", "author", "\"First\"", "'Second'", "First"), + ("author", "author", "Same", "Same", "Same"), + ("author", "author", "''", "Second", "''"), + ("author", "author", "\"\"", "Second", "\"\""), + ("author", "author", "\" First \"", "Second", " First "), + ("author", "author", "0", "Second", "0"), + ("description", "description", "First", "Second", "First"), + ("metadata", "metadata", "First", "Second", "First"), + ("vendor-key", "vendor-key", "First", "Second", "First"), + ("vendor_key", "vendor_key", "First", "Second", "First"), + ("author", "Author", "First", "Second", "First"), + ("Author", "author", "First", "Second", "First"), + ("author", "AUTHOR", "First", "Second", "First"), + ("description", "Description", "First", "Second", "First"), + ("metadata", "Metadata", "First", "Second", "First"), + ("vendor-key", "Vendor-Key", "First", "Second", "First"), + ("author", "author", "\n 'First'", "Second", "First"), + ("author", "Author", "First", "\n \"Second\"", "First"), + ("author", "author", "\"author: First\"", "Second", "author: First"), + ("author", "author", "First\n # author: Not a field", "Second", "First"), + ("author", "author", "|-\n First\n paragraph", "Second", "|-"), + ("author", "Author", "First", ">-\n Second\n paragraph", "First"), + ("author", "author", ">\n First\n paragraph", "|\n Second\n paragraph", ">"), + ]; + foreach (var (key, duplicateKey, value, duplicateValue, expected) in cases) + { + data.Add(key, duplicateKey, value, duplicateValue, expected, "\n"); + data.Add(key, duplicateKey, value, duplicateValue, expected, "\r\n"); + } + + return data; + } + } + + [Theory] + [MemberData(nameof(ExistingScalarFormats))] + public async Task GetSkillsAsync_ExistingScalarFormats_PreservedAsync( + string value, string expected, string newline, bool metadataFirst, bool duplicateMetadata, string quote) + { + // Arrange + string fields = $"{quote}name{quote}: test-skill\n{quote}description{quote}: {value}\n" + + $"{quote}license{quote}: 'MIT'\n{quote}compatibility{quote}: Any runtime\n{quote}allowed-tools{quote}: read\n"; + string metadata = $"{quote}metadata{quote}:\n author: test\n" + (duplicateMetadata ? " author: Ignored\n" : ""); + string content = "---\n" + (metadataFirst ? metadata + fields : fields + metadata) + "---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal(expected, frontmatter.Description); + Assert.Equal("MIT", frontmatter.License); + Assert.Equal("Any runtime", frontmatter.Compatibility); + Assert.Equal("read", frontmatter.AllowedTools); + Assert.NotNull(frontmatter.Metadata); + Assert.Equal("test", frontmatter.Metadata["author"]); + } + + [Theory] + [MemberData(nameof(ExistingMetadataScalarFormats))] + public async Task GetSkillsAsync_ExistingMetadataScalarFormats_PreservedAsync(string value, string expected, string newline) + { + // Arrange + string content = + "---\nname: test-skill\ndescription: Read files\n" + + $"metadata:\n author: {value}\n version: '1.0'\nlicense: MIT\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert - preserve existing scalar parsing rather than expanding YAML support. + var skill = Assert.Single(skills); + Assert.Equal("Read files", skill.Frontmatter.Description); + Assert.Equal("MIT", skill.Frontmatter.License); + Assert.NotNull(skill.Frontmatter.Metadata); + Assert.Equal(2, skill.Frontmatter.Metadata.Count); + Assert.Equal(expected, skill.Frontmatter.Metadata["author"]); + Assert.Equal("1.0", skill.Frontmatter.Metadata["version"]); + Assert.Contains("Body.", await skill.GetContentAsync()); + logger.Verify( + l => l.Log( + It.Is(level => level >= LogLevel.Warning), + It.IsAny(), + It.IsAny(), + It.IsAny(), + It.IsAny>()), + Times.Never); + } + + [Theory] + [MemberData(nameof(ConflictingScalarFormats))] + public async Task GetSkillsAsync_ConflictingScalarFormats_ExcludesSkillAndLogsErrorAsync( + string firstValue, + string secondValue, + string secondKey) + { + // Arrange + string content = $"---\nname: test-skill\ndescription: {firstValue}\n{secondKey}: {secondValue}\nlicense: MIT\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + string diagnostic = secondKey == "description" + ? "duplicate frontmatter field 'description'" + : $"incorrectly cased frontmatter field '{secondKey}'; expected 'description'"; + logger.Verify( + l => l.Log( + LogLevel.Error, + It.IsAny(), + It.Is((state, _) => state.ToString()!.IndexOf(diagnostic, StringComparison.Ordinal) >= 0), + It.IsAny(), + It.IsAny>()), + Times.Once); + } + + [Fact] + public async Task GetSkillsAsync_NestedMetadataAndUnknownFields_KeepsFirstMetadataValueAsync() + { + // Arrange + const string Content = + "---\nname: test-skill\ndescription: Read files\n" + + "# description: Not a field\n" + + "metadata:\n author: First\n author: Second\n Description: Nested text\n" + + " \"description\": Ignored\n 'metadata': Ignored\n" + + "vendor-option: First\nvendor-option: Second\n" + + "'vendor-option': Third\n\"VENDOR-OPTION\": Fourth\n---\n\"description\": Body text"; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", Content); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.NotNull(frontmatter.Metadata); + Assert.Equal(2, frontmatter.Metadata.Count); + Assert.Equal("First", frontmatter.Metadata["author"]); + Assert.Equal("Nested text", frontmatter.Metadata["Description"]); + } + + [Theory] + [MemberData(nameof(DuplicateMetadataEntries))] + public async Task GetSkillsAsync_DuplicateMetadata_KeepsFirstValueAndLogsWarningAsync( + string key, + string duplicateKey, + string value, + string duplicateValue, + string expected, + string newline) + { + // Arrange + string content = + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + + $" {key}: {value}\n other: Preserved\n {duplicateKey}: {duplicateValue}\n" + + $" {duplicateKey}: Third\n tail: Retained\nlicense: MIT\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var skill = Assert.Single(skills); + var frontmatter = skill.Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Equal("MIT", frontmatter.License); + Assert.Contains("Body.", await skill.GetContentAsync()); + Assert.NotNull(frontmatter.Metadata); + Assert.Equal(3, frontmatter.Metadata.Count); + Assert.Equal(expected, frontmatter.Metadata[key]); + Assert.Equal(expected, frontmatter.Metadata[duplicateKey]); + Assert.Equal("Preserved", frontmatter.Metadata["other"]); + Assert.Equal("Retained", frontmatter.Metadata["tail"]); + Assert.Equal(key, Assert.Single(frontmatter.Metadata.Keys, k => string.Equals(k, key, StringComparison.OrdinalIgnoreCase))); + string diagnostic = $"duplicate metadata key '{duplicateKey}'; keeping the first value"; + logger.Verify( + l => l.Log( + LogLevel.Warning, + It.IsAny(), + It.Is((state, _) => + state.ToString()!.IndexOf("SKILL.md", StringComparison.Ordinal) >= 0 && + state.ToString()!.IndexOf(diagnostic, StringComparison.Ordinal) >= 0), + It.IsAny(), + It.IsAny>()), + Times.Exactly(2)); + } + + [Theory] + [InlineData("''")] + [InlineData("\"\"")] + public async Task GetSkillsAsync_EmptyQuotedOptionalValues_KeepExistingRepresentationAsync(string value) + { + // Arrange + string content = + "---\nname: test-skill\ndescription: Read files\n" + + $"license: {value}\ncompatibility: {value}\nallowed-tools: {value}\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert - empty-quote normalization belongs to the separate YAML conformance work. + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal(value, frontmatter.License); + Assert.Equal(value, frontmatter.Compatibility); + Assert.Equal(value, frontmatter.AllowedTools); + } + + [Theory] + [InlineData("|")] + [InlineData("|-")] + [InlineData("|+")] + [InlineData(">")] + [InlineData(">-")] + [InlineData(">+")] + public async Task GetSkillsAsync_EmptyBlockBeforeAnotherField_ExcludesSkillAsync(string value) + { + // Arrange + _ = this.CreateSkillDirectoryWithRawContent( + "test-skill", $"---\nname: test-skill\ndescription: {value}\n\nlicense: MIT\n---\nBody."); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + [Fact] public async Task GetSkillsAsync_ValidSkill_ReturnsSkillAsync() { @@ -192,6 +634,211 @@ public async Task GetSkillsAsync_MissingDescriptionField_ExcludesSkillAsync() Assert.Empty(skills); } + [Theory] + [InlineData("name")] + [InlineData("description")] + [InlineData("license")] + [InlineData("compatibility")] + [InlineData("metadata")] + [InlineData("allowed-tools")] + public async Task GetSkillsAsync_DuplicateRecognizedFrontmatterField_ExcludesSkillAsync(string field) + { + // Arrange + string fields = field switch + { + "name" => "name: ambiguous-skill\nname: ambiguous-skill\ndescription: A skill", + "description" => "name: ambiguous-skill\ndescription: A skill\ndescription: A second description", + "metadata" => "name: ambiguous-skill\ndescription: A skill\nmetadata:\n author: first\nmetadata:\n version: 1.0", + _ => $"name: ambiguous-skill\ndescription: A skill\n{field}: first\n{field}: second", + }; + _ = this.CreateSkillDirectoryWithRawContent( + "ambiguous-skill", + $"---\n{fields}\n---\nBody."); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + + [Theory] + [InlineData("name", "Name")] + [InlineData("description", "Description")] + [InlineData("license", "License")] + [InlineData("compatibility", "Compatibility")] + [InlineData("metadata", "Metadata")] + [InlineData("allowed-tools", "Allowed-Tools")] + public async Task GetSkillsAsync_IncorrectlyCasedRecognizedFrontmatterField_ExcludesSkillAsync( + string field, + string incorrectlyCasedField) + { + // Arrange + string fields = field switch + { + "name" => $"{incorrectlyCasedField}: ambiguous-skill\ndescription: A skill", + "description" => $"name: ambiguous-skill\n{incorrectlyCasedField}: A skill", + "metadata" => $"name: ambiguous-skill\ndescription: A skill\n{incorrectlyCasedField}:\n author: test", + _ => $"name: ambiguous-skill\ndescription: A skill\n{incorrectlyCasedField}: value", + }; + _ = this.CreateSkillDirectoryWithRawContent( + "ambiguous-skill", + $"---\n{fields}\n---\nBody."); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + + [Theory] + [MemberData(nameof(ConflictingQuotedFields))] + public async Task GetSkillsAsync_ConflictingQuotedFields_ExcludesSkillAndLogsErrorAsync( + string content, + string diagnostic) + { + // Arrange + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content); + var logger = new Mock(); + logger.Setup(l => l.IsEnabled(It.IsAny())).Returns(true); + var loggerFactory = new Mock(); + loggerFactory.Setup(f => f.CreateLogger(It.IsAny())).Returns(logger.Object); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor, loggerFactory: loggerFactory.Object); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + logger.Verify( + l => l.Log( + LogLevel.Error, + It.IsAny(), + It.Is((state, _) => state.ToString()!.IndexOf(diagnostic, StringComparison.Ordinal) >= 0), + It.IsAny(), + It.IsAny>()), + Times.Once); + } + + [Theory] + [InlineData("'", "\n")] + [InlineData("'", "\r\n")] + [InlineData("\"", "\n")] + [InlineData("\"", "\r\n")] + public async Task GetSkillsAsync_QuotedEmptyOptionalFields_RemainUnsetAsync(string quote, string newline) + { + // Arrange + string content = + $"---\n{quote}metadata{quote}:\n{quote}name{quote}: test-skill\n" + + $"{quote}description{quote}: Read files\n{quote}license{quote}:\n" + + $"{quote}compatibility{quote}:\n{quote}allowed-tools{quote}:\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("test-skill", frontmatter.Name); + Assert.Equal("Read files", frontmatter.Description); + Assert.Null(frontmatter.Metadata); + Assert.Null(frontmatter.License); + Assert.Null(frontmatter.Compatibility); + Assert.Null(frontmatter.AllowedTools); + } + + [Theory] + [InlineData("metadata", "\n")] + [InlineData("metadata", "\r\n")] + [InlineData("license", "\n")] + [InlineData("license", "\r\n")] + [InlineData("compatibility", "\n")] + [InlineData("compatibility", "\r\n")] + [InlineData("allowed-tools", "\n")] + [InlineData("allowed-tools", "\r\n")] + [InlineData("vendor-option", "\n")] + [InlineData("vendor-option", "\r\n")] + public async Task GetSkillsAsync_EmptyInlineValue_DoesNotConsumeNextFieldAsync(string field, string newline) + { + // Arrange + string content = $"---\n{field}: \t\nname: test-skill\ndescription: A test skill.\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var skill = Assert.Single(skills); + Assert.Equal("test-skill", skill.Frontmatter.Name); + Assert.Equal("A test skill.", skill.Frontmatter.Description); + Assert.Null(skill.Frontmatter.License); + Assert.Null(skill.Frontmatter.Compatibility); + Assert.Null(skill.Frontmatter.AllowedTools); + } + + [Theory] + [InlineData("license", "\n")] + [InlineData("license", "\r\n")] + [InlineData("compatibility", "\n")] + [InlineData("compatibility", "\r\n")] + [InlineData("allowed-tools", "\n")] + [InlineData("allowed-tools", "\r\n")] + public async Task GetSkillsAsync_EmptyOptionalScalarAtEnd_RemainsNullAsync(string field, string newline) + { + // Arrange + string content = $"---\nname: test-skill\ndescription: Read files\n{field}: \t\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + var frontmatter = Assert.Single(skills).Frontmatter; + Assert.Equal("Read files", frontmatter.Description); + Assert.Null(frontmatter.License); + Assert.Null(frontmatter.Compatibility); + Assert.Null(frontmatter.AllowedTools); + } + + [Theory] + [InlineData("name:", "\n")] + [InlineData("name:", "\r\n")] + [InlineData("description:", "\n")] + [InlineData("description:", "\r\n")] + [InlineData("metadata:\nmetadata:", "\n")] + [InlineData("metadata:\nmetadata:", "\r\n")] + [InlineData("license:\nlicense: MIT", "\n")] + [InlineData("license:\nlicense: MIT", "\r\n")] + [InlineData("license: MIT\nlicense:", "\n")] + [InlineData("license: MIT\nlicense:", "\r\n")] + [InlineData("compatibility:\ncompatibility: Any runtime", "\n")] + [InlineData("compatibility:\ncompatibility: Any runtime", "\r\n")] + [InlineData("allowed-tools:\nallowed-tools: read", "\n")] + [InlineData("allowed-tools:\nallowed-tools: read", "\r\n")] + [InlineData("Metadata:", "\n")] + [InlineData("Metadata:", "\r\n")] + [InlineData("allowed-tools: read\nALLOWED-TOOLS:", "\n")] + [InlineData("allowed-tools: read\nALLOWED-TOOLS:", "\r\n")] + public async Task GetSkillsAsync_EmptyInlineValue_DoesNotBypassKeyValidationAsync(string fields, string newline) + { + // Arrange + string content = $"---\n{fields}\nname: test-skill\ndescription: A test skill.\n---\nBody."; + _ = this.CreateSkillDirectoryWithRawContent("test-skill", content.Replace("\n", newline)); + var source = new AgentFileSkillsSource(this._testRoot, s_noOpExecutor); + + // Act + var skills = await source.GetSkillsAsync(TestAgentSkillsSourceContextFactory.Create()); + + // Assert + Assert.Empty(skills); + } + [Theory] [InlineData("BadName")] [InlineData("-leading-hyphen")] diff --git a/python/packages/ag-ui/AGENTS.md b/python/packages/ag-ui/AGENTS.md index 25233eaf56d..94a11e1ca29 100644 --- a/python/packages/ag-ui/AGENTS.md +++ b/python/packages/ag-ui/AGENTS.md @@ -41,6 +41,10 @@ AG-UI protocol integration for building agent UIs with the AG-UI standard. `Agent.run` path, including Agent middleware, provider `before_run` preparation, and function middleware. `_approval_execution.py` tracks execution through existing function/chat middleware interfaces; it does not prepare providers or invoke the tools itself. `MiddlewareFailure` remains fatal. +- Programmatic `AgentFrameworkAgent.run(...)` and `run_agent_stream(...)` calls accept + `function_invocation_kwargs` and preserve that tool-runtime context across approval resume. Adapter-local execution + merges only that mapping, the wrapped agent's default `additional_function_arguments`, and the current run options + (in increasing precedence), not unrelated AG-UI run kwargs. - Queued approvals follow the wrapped Agent's scheduling. With `ToolApprovalMiddleware`, collected decisions remain pending until that middleware releases the batch; AG-UI does not execute a tool ahead of it. Collected server-side grants do not require the client to submit the same approval again. diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py index 9d127f2bbff..ab179da165c 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent.py @@ -3,7 +3,7 @@ """AgentFrameworkAgent wrapper for AG-UI protocol.""" import warnings -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Mapping from typing import Any, cast from ag_ui.core import BaseEvent @@ -175,11 +175,14 @@ def snapshot_store(self) -> AGUIThreadSnapshotStore | None: async def run( self, input_data: dict[str, Any], + *, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AsyncGenerator[BaseEvent, None]: """Run the wrapped agent and yield AG-UI events. Args: input_data: The AG-UI run input containing messages, state, etc. + function_invocation_kwargs: Keyword arguments forwarded only to tool invocation. Yields: AG-UI events @@ -190,5 +193,6 @@ async def run( self.agent, self.config, approval_state_store=self._approval_state_store, + function_invocation_kwargs=function_invocation_kwargs, ): yield event diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py index 825a60ec621..1baf001c640 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_agent_run.py @@ -1938,7 +1938,22 @@ async def forward_hosted_decision(approval: Content = approval) -> list[Content] raise TypeError("Local approval execution with context providers requires a core Agent execution path.") client = getattr(agent, "client", None) config = normalize_function_invocation_configuration(getattr(client, "function_invocation_configuration", None)) - tool_kwargs = {k: v for k, v in run_kwargs.items() if k != "options"} + raw_function_invocation_kwargs = run_kwargs.get("function_invocation_kwargs") + tool_kwargs = ( + dict(cast(Mapping[str, Any], raw_function_invocation_kwargs)) + if raw_function_invocation_kwargs is not None + else {} + ) + default_options = getattr(agent, "default_options", None) + if isinstance(default_options, Mapping) and ( + default_function_arguments := default_options.get("additional_function_arguments") + ): + tool_kwargs.update(cast(Mapping[str, Any], default_function_arguments)) + raw_options = run_kwargs.get("options") + if isinstance(raw_options, Mapping) and ( + additional_function_arguments := raw_options.get("additional_function_arguments") + ): + tool_kwargs.update(cast(Mapping[str, Any], additional_function_arguments)) for approval in static_approved: function_call = approval.function_call call_id = (function_call.call_id if function_call else None) or approval.id or "" @@ -2668,6 +2683,8 @@ async def run_agent_stream( agent: SupportsAgentRun, config: AgentConfig, approval_state_store: InMemoryAGUIApprovalStateStore | None = None, + *, + function_invocation_kwargs: Mapping[str, Any] | None = None, ) -> AsyncGenerator[BaseEvent]: """Run agent and yield AG-UI events. @@ -2680,6 +2697,7 @@ async def run_agent_stream( config: Agent configuration approval_state_store: Optional server-side Approval State store used to preserve approval-only middleware state across AG-UI requests. + function_invocation_kwargs: Keyword arguments forwarded only to tool invocation. Yields: AG-UI events @@ -2692,6 +2710,7 @@ async def run_agent_stream( config, state_store, authorized_executions=authorized_executions, + function_invocation_kwargs=function_invocation_kwargs, ) try: async for event in stream: @@ -2711,6 +2730,7 @@ async def _run_agent_stream( approval_state_store: InMemoryAGUIApprovalStateStore, *, authorized_executions: dict[ApprovalOccurrenceIdentity, AuthorizedExecution], + function_invocation_kwargs: Mapping[str, Any] | None, ) -> AsyncGenerator[BaseEvent]: # Parse IDs supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") @@ -3105,6 +3125,8 @@ async def _run_agent_stream( # Build run kwargs (Feature #6: Azure store flag when metadata present) run_kwargs: dict[str, Any] = {"session": session} + if function_invocation_kwargs is not None: + run_kwargs["function_invocation_kwargs"] = dict(function_invocation_kwargs) if tools: run_kwargs["tools"] = tools # Hand the forwarded AG-UI context to the A2UI runner PER REQUEST (not just at diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py index 539ddd21855..33d3bede30e 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow.py @@ -525,6 +525,12 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: (``forwarded_props: {"checkpoint_id": ...}``), which restores the persisted workflow state instead of starting a fresh turn. + Resume also works when storage lives only on the resolved workflow + (``WorkflowBuilder(checkpoint_storage=...)``) and this wrapper / endpoint omits + ``checkpoint_storage``, so emitted pause ids remain round-trippable for hosts that + expose the workflow through ``AgentFrameworkWorkflow`` or + ``add_agent_framework_fastapi_endpoint`` without duplicating storage configuration. + Note: Checkpointing (the ``agent_framework`` workflow checkpoint mechanism) is independent from AG-UI Thread Snapshot persistence (``snapshot_store``). @@ -547,11 +553,6 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: checkpoint_storage = self.checkpoint_storage checkpoint_id = _checkpoint_id_from_input(input_data) - if checkpoint_id is not None and checkpoint_storage is None: - raise ValueError( - "Resuming from a checkpoint requires checkpoint_storage to be configured on " - "AgentFrameworkWorkflow (or the AG-UI endpoint)." - ) supplied_thread_id = input_data.get("thread_id") or input_data.get("threadId") request_owner = (snapshot_scope, str(supplied_thread_id) if supplied_thread_id is not None else None) @@ -561,6 +562,15 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: if interrupt.get("id") is not None } workflow = self._resolve_workflow(thread_id, snapshot_scope) + # Prefer wrapper/endpoint storage; otherwise allow builder/runtime storage on the + # resolved workflow so pause ids emitted without AG-UI storage still resume. + if checkpoint_id is not None and checkpoint_storage is None: + if not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] + raise ValueError( + "Resuming from a checkpoint requires checkpoint_storage to be configured on " + "AgentFrameworkWorkflow (or the AG-UI endpoint), or WorkflowBuilder " + "checkpoint storage on the workflow instance." + ) live_pending_events = await _pending_request_events(self.workflow) if self.workflow is not None else {} if self.workflow is not None and checkpoint_id is None: for request_event in live_pending_events.values(): @@ -572,9 +582,14 @@ async def run(self, input_data: dict[str, Any]) -> AsyncGenerator[BaseEvent]: code="WORKFLOW_RESUME_NOT_FOUND", ) return - if checkpoint_id is not None and checkpoint_storage is not None: + if checkpoint_id is not None: try: - checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + else: + checkpoint = await workflow._runner.context.load_checkpoint(checkpoint_id) # pyright: ignore[reportPrivateUsage] + if checkpoint is None: + raise LookupError(f"checkpoint '{checkpoint_id}' was not found") except Exception as exc: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) yield RunErrorEvent( diff --git a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py index 9ee70075f2b..9a5126cc813 100644 --- a/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py +++ b/python/packages/ag-ui/agent_framework_ag_ui/_workflow_run.py @@ -8,7 +8,7 @@ import json import logging import uuid -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Awaitable, Callable from functools import partial from types import UnionType from typing import Any, Union, cast, get_args, get_origin, get_type_hints @@ -61,6 +61,8 @@ logger = logging.getLogger(__name__) +_BASELINE_OMITTED = object() + _PUBLIC_WORKFLOW_ERROR_MESSAGE = "Workflow execution failed." @@ -130,7 +132,10 @@ def _workflow_interrupt_value(request_data: Any) -> Any: return {"data": safe_request_data} -def _workflow_interrupt_metadata(request_payload: dict[str, Any], value: Any) -> dict[str, Any]: +def _workflow_interrupt_metadata( + request_payload: dict[str, Any], + value: Any, +) -> dict[str, Any]: """Build Agent Framework metadata for workflow request_info interrupts.""" agent_framework_metadata = { key: make_json_safe(value) @@ -148,6 +153,99 @@ def _workflow_interrupt_metadata(request_payload: dict[str, Any], value: Any) -> return {"agent_framework": agent_framework_metadata} +def _attach_checkpoint_id_to_interrupts( + interrupts: list[dict[str, Any]], + checkpoint_id: str | None, +) -> list[dict[str, Any]]: + """Attach ``checkpoint_id`` to each interrupt's ``metadata.agent_framework``. + + Multi-worker hosts need the pause checkpoint on the wire so the next resume can pass + ``forwardedProps.checkpoint_id`` without a side-channel lookup. No-op when checkpointing + is inactive or the id is already present. + """ + if not checkpoint_id or not interrupts: + return interrupts + + attached: list[dict[str, Any]] = [] + for interrupt in interrupts: + entry = dict(interrupt) + metadata = entry.get("metadata") + if isinstance(metadata, dict): + metadata = dict(metadata) + else: + metadata = {} + agent_framework = metadata.get("agent_framework") + if isinstance(agent_framework, dict): + agent_framework = dict(agent_framework) + else: + agent_framework = {} + agent_framework.setdefault("checkpoint_id", checkpoint_id) + metadata["agent_framework"] = agent_framework + entry["metadata"] = metadata + attached.append(entry) + return attached + + +def _interrupt_request_ids(interrupts: list[dict[str, Any]]) -> set[str]: + return {str(item["id"]) for item in interrupts if item.get("id") is not None} + + +async def _pause_checkpoint_id_for_interrupts( + *, + workflow: Workflow, + checkpoint_storage: CheckpointStorage | None, + interrupts: list[dict[str, Any]], + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: Any = _BASELINE_OMITTED, +) -> str | None: + """Resolve the pause checkpoint for *this* run's interrupts via core. + + When ``baseline_checkpoint_id`` is omitted, core uses the baseline captured at + ``workflow.run()`` start. Pass ``None`` explicitly for short-circuit paths that + did not call ``run()`` and should advertise the current runner id when present. + """ + if not interrupts: + return None + + resolve = getattr(workflow, "resolve_pause_checkpoint_id", None) + if not callable(resolve): + return None + + # getattr returns a plain object to the type checker; cast to an awaitable callable. + resolve_fn = cast( + Callable[..., Awaitable[str | None]], + resolve, + ) + kwargs: dict[str, Any] = { + "checkpoint_storage": checkpoint_storage, + "known_checkpoint_id": known_checkpoint_id, + } + if baseline_checkpoint_id is not _BASELINE_OMITTED: + kwargs["baseline_checkpoint_id"] = baseline_checkpoint_id + return await resolve_fn(_interrupt_request_ids(interrupts), **kwargs) + + +async def _interrupts_with_pause_checkpoint( + *, + interrupts: list[dict[str, Any]], + workflow: Workflow, + checkpoint_storage: CheckpointStorage | None, + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: Any = _BASELINE_OMITTED, +) -> list[dict[str, Any]]: + """Attach a run-scoped pause checkpoint id to interrupts when available.""" + if not interrupts: + return interrupts + pause_checkpoint_id = await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=checkpoint_storage, + interrupts=interrupts, + known_checkpoint_id=known_checkpoint_id, + baseline_checkpoint_id=baseline_checkpoint_id, + ) + return _attach_checkpoint_id_to_interrupts(interrupts, pause_checkpoint_id) + + async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: """Best-effort retrieval of pending request_info events from workflow context.""" runner_context = getattr(workflow, "_runner_context", None) @@ -171,7 +269,9 @@ async def _pending_request_events(workflow: Workflow) -> dict[str, Any]: async def _pending_request_events_from_checkpoint( checkpoint_id: str, - checkpoint_storage: CheckpointStorage, + checkpoint_storage: CheckpointStorage | None = None, + *, + workflow: Any | None = None, ) -> dict[str, Any]: """Read pending request_info events from a persisted checkpoint without restoring it. @@ -182,16 +282,34 @@ async def _pending_request_events_from_checkpoint( exposes them without running any executor ``on_checkpoint_restore`` hook; the single ``workflow.run(checkpoint_id=...)`` then performs the one real restore, so the restore -- and every custom restore hook -- runs exactly once per resume. + + ``checkpoint_storage`` is preferred when provided. Otherwise the workflow's + effective builder/runtime storage is used when ``has_checkpointing()`` is true, + so AG-UI can round-trip pause IDs emitted from ``WorkflowBuilder(checkpoint_storage=...)`` + without requiring a duplicate AG-UI storage argument. """ try: - checkpoint = await checkpoint_storage.load(checkpoint_id) + if checkpoint_storage is not None: + checkpoint = await checkpoint_storage.load(checkpoint_id) + else: + context = getattr(getattr(workflow, "_runner", None), "context", None) + if context is None or not context.has_checkpointing(): + raise ValueError( + "Resuming a checkpoint with an AG-UI resume payload requires checkpoint_storage " + "(or WorkflowBuilder checkpoint storage on the workflow instance)." + ) + checkpoint = await context.load_checkpoint(checkpoint_id) + except ValueError: + raise except Exception: logger.warning( "Could not load checkpoint for resume-response coercion; the core run will surface any error.", exc_info=True, ) return {} - return dict(checkpoint.pending_request_info_events) + if checkpoint is None: + return {} + return dict(checkpoint.pending_request_info_events or {}) def _interrupt_entry_for_request_event(request_event: Any) -> dict[str, Any] | None: @@ -1062,9 +1180,15 @@ async def run_workflow_stream( # pure checkpoint restore still surfaces its pending interrupts instead of tripping # the "resume required" contract. if checkpoint_id is not None and resume_payload is not None: - if checkpoint_storage is None: + # Prefer the explicit AG-UI storage argument; otherwise allow the workflow's + # builder/runtime storage so builder-emitted pause IDs remain round-trippable. + if checkpoint_storage is None and not workflow._runner.context.has_checkpointing(): # pyright: ignore[reportPrivateUsage] raise ValueError("Resuming a checkpoint with an AG-UI resume payload requires checkpoint_storage.") - pending_before_run = await _pending_request_events_from_checkpoint(checkpoint_id, checkpoint_storage) + pending_before_run = await _pending_request_events_from_checkpoint( + checkpoint_id, + checkpoint_storage, + workflow=workflow, + ) else: pending_before_run = await _pending_request_events(workflow) pending_interrupt_ids = _pending_workflow_interrupt_ids(pending_before_run) @@ -1140,12 +1264,30 @@ async def run_workflow_stream( interrupt_event_value = _workflow_interrupt_event_value(request_payload) if interrupt_event_value is not None: yield CustomEvent(name=_INTERRUPT_CARD_EVENT_NAME, value=interrupt_event_value) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=pending_interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + baseline_checkpoint_id=None, + ), + ) return if checkpoint_id is None and not responses and not messages: yield RunStartedEvent(run_id=run_id, thread_id=thread_id) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=pending_interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=pending_interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + baseline_checkpoint_id=None, + ), + ) return def _drain_open_message() -> list[TextMessageEndEvent]: @@ -1281,7 +1423,15 @@ def _drain_open_blocks() -> list[BaseEvent]: yield end_event if not interrupts: interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow))) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + ), + ) terminal_emitted = True elif state_value not in _TERMINAL_STATES: yield CustomEvent(name="status", value={"state": state_value}) @@ -1447,4 +1597,12 @@ def _drain_open_blocks() -> list[BaseEvent]: if not terminal_emitted and not run_error_emitted: if not interrupts: interrupts.extend(_interrupts_from_pending_requests(await _pending_request_events(workflow))) - yield _build_run_finished_event(run_id=run_id, thread_id=thread_id, interrupts=interrupts) + yield _build_run_finished_event( + run_id=run_id, + thread_id=thread_id, + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupts, + workflow=workflow, + checkpoint_storage=checkpoint_storage, + ), + ) diff --git a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py index b32cea9af8c..2a01c03a4b4 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py +++ b/python/packages/ag-ui/tests/ag_ui/test_agent_wrapper_comprehensive.py @@ -8,7 +8,16 @@ import pytest from ag_ui.core import RunErrorEvent, ToolCallResultEvent -from agent_framework import Agent, ChatOptions, ChatResponseUpdate, Content, Message +from agent_framework import ( + Agent, + ChatOptions, + ChatResponseUpdate, + Content, + ContextProvider, + FunctionInvocationContext, + Message, + tool, +) from pydantic import BaseModel from agent_framework_ag_ui._approval_lifecycle import ApprovalExecutionOwner @@ -23,6 +32,82 @@ def _approval_request_id(events: list[Any]) -> str: return str(event.value["id"]) +async def test_approval_resume_with_context_provider_preserves_function_invocation_kwargs( + streaming_chat_client_stub: Any, +) -> None: + """Context-provider preparation and runtime tool kwargs survive an AG-UI resume.""" + from agent_framework.ag_ui import AgentFrameworkAgent + + observed_kwargs: list[dict[str, Any]] = [] + phase = "pause" + + class NoOpContextProvider(ContextProvider): + pass + + @tool(name="sensitive_action", approval_mode="always_require") + def sensitive_action(context: FunctionInvocationContext) -> str: + observed_kwargs.append(dict(context.kwargs)) + return "executed" + + async def stream_fn( + messages: MutableSequence[Message], + options: ChatOptions, + **kwargs: Any, + ) -> AsyncIterator[ChatResponseUpdate]: + del messages, options, kwargs + if phase == "pause": + yield ChatResponseUpdate( + contents=[ + Content.from_function_call( + name="sensitive_action", + call_id="call-sensitive", + arguments="{}", + ) + ] + ) + return + yield ChatResponseUpdate(contents=[Content.from_text(text="Done")]) + + wrapper = AgentFrameworkAgent( + agent=Agent( + client=streaming_chat_client_stub(stream_fn), + name="test_agent", + instructions="Test", + tools=[sensitive_action], + context_providers=[NoOpContextProvider("context")], + ) + ) + runtime_kwargs = {"user_id": "user-123"} + thread_id = "context-provider-runtime-kwargs" + pause_events = [ + event + async for event in wrapper.run( + {"thread_id": thread_id, "messages": [{"role": "user", "content": "do it"}]}, + function_invocation_kwargs=runtime_kwargs, + ) + ] + approval_id = _approval_request_id(pause_events) + + phase = "resume" + resume_events = [ + event + async for event in wrapper.run( + { + "thread_id": thread_id, + "messages": [], + "resume": [{"interruptId": approval_id, "status": "resolved", "payload": {"accepted": True}}], + }, + function_invocation_kwargs=runtime_kwargs, + ) + ] + + assert observed_kwargs[0]["user_id"] == runtime_kwargs["user_id"] + assert set(observed_kwargs[0]) == {"session", "user_id"} + assert [ + (event.tool_call_id, event.content) for event in resume_events if isinstance(event, ToolCallResultEvent) + ] == [("call-sensitive", "executed")] + + async def test_agent_initialization_basic(streaming_chat_client_stub): """Test basic agent initialization without state schema.""" from agent_framework.ag_ui import AgentFrameworkAgent diff --git a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py index a43dde00d19..86b52f068aa 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py +++ b/python/packages/ag-ui/tests/ag_ui/test_approval_result_event.py @@ -8,7 +8,7 @@ from typing import Any from ag_ui.core import ToolCallResultEvent -from agent_framework import AgentResponseUpdate, Content, FunctionTool +from agent_framework import AgentResponseUpdate, Content, FunctionInvocationContext, FunctionTool from agent_framework.exceptions import UserInputRequiredException from conftest import StubAgent # pyrefly: ignore[missing-import] # pyright: ignore[reportMissingImports] @@ -88,10 +88,18 @@ async def _run_resume( return events -async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent]: +async def _run_custom_approval( + tool: FunctionTool, + *, + function_invocation_kwargs: dict[str, Any] | None = None, + default_additional_function_arguments: dict[str, Any] | None = None, +) -> tuple[list[Any], StubAgent]: agent = StubAgent( updates=[AgentResponseUpdate(contents=[Content.from_text(text="Done.")], role="assistant")], - default_options={"tools": [tool]}, + default_options={ + "tools": [tool], + "additional_function_arguments": default_additional_function_arguments, + }, ) store = InMemoryAGUIApprovalStateStore() store.register( @@ -114,6 +122,7 @@ async def _run_custom_approval(tool: FunctionTool) -> tuple[list[Any], StubAgent agent, AgentConfig(), approval_state_store=store, + function_invocation_kwargs=function_invocation_kwargs, ) ] return events, agent @@ -134,6 +143,30 @@ async def test_approved_call_emits_one_live_result_under_original_identity() -> assert [(event.tool_call_id, event.content) for event in results] == [("call-weather", "Sunny in Seattle")] +async def test_approved_call_merges_function_invocation_kwargs_with_agent_defaults() -> None: + """Static approval execution uses normal runtime/default precedence without run-kwarg leakage.""" + observed_kwargs: list[dict[str, Any]] = [] + + def inspect_context(context: FunctionInvocationContext) -> str: + observed_kwargs.append(dict(context.kwargs)) + return "inspected" + + events, _ = await _run_custom_approval( + FunctionTool(name="inspect_context", description="Inspect context", func=inspect_context), + function_invocation_kwargs={"runtime_only": "runtime", "shared": "runtime"}, + default_additional_function_arguments={"default_only": "default", "shared": "default"}, + ) + + results = [event for event in events if isinstance(event, ToolCallResultEvent)] + assert observed_kwargs[0] == { + "default_only": "default", + "runtime_only": "runtime", + "session": observed_kwargs[0]["session"], + "shared": "default", + } + assert [(event.tool_call_id, event.content) for event in results] == [("call-custom", "inspected")] + + async def test_rejected_call_does_not_execute_or_emit_live_result() -> None: executions: list[str] = [] diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py index 8702adcd047..615d3ce4f1a 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_agent.py @@ -260,7 +260,7 @@ async def test_workflow_run_reads_checkpoint_id_from_camelcase_forwarded_props() async def test_workflow_resume_without_checkpoint_storage_raises() -> None: - """Requesting a checkpoint resume without configured storage should fail loudly.""" + """Requesting a checkpoint resume without any storage should fail loudly.""" workflow = _build_multi_superstep_workflow() agent = AgentFrameworkWorkflow(workflow=workflow) @@ -275,6 +275,39 @@ async def test_workflow_resume_without_checkpoint_storage_raises() -> None: ) +async def test_workflow_wrapper_resumes_builder_storage_without_agui_storage() -> None: + """Builder-owned storage must round-trip through AgentFrameworkWorkflow.run() without wrapper storage.""" + storage = InMemoryCheckpointStorage() + workflow = _build_multi_superstep_workflow(storage) + # Host configures storage only on the builder; AG-UI wrapper / endpoint omit it. + agent = AgentFrameworkWorkflow(workflow=workflow) + + first_events = await _run( + agent, + {"thread_id": "thread-builder-cp", "messages": [{"role": "user", "content": "start"}]}, + ) + assert "RUN_ERROR" not in [event.type for event in first_events] + + checkpoints = sorted( + await storage.list_checkpoints(workflow_name=workflow.name), + key=lambda checkpoint: checkpoint.timestamp, + ) + assert checkpoints, "expected the builder-storage run to create at least one checkpoint" + resume_checkpoint_id = checkpoints[0].checkpoint_id + + resume_events = await _run( + agent, + { + "thread_id": "thread-builder-cp", + "messages": [], + "forwarded_props": {"checkpoint_id": resume_checkpoint_id}, + }, + ) + resumed_types = [event.type for event in resume_events] + assert "RUN_FINISHED" in resumed_types + assert "RUN_ERROR" not in resumed_types + + async def test_workflow_run_without_checkpointing_is_unchanged() -> None: """Existing run(input_data) calls keep working unchanged when no checkpoint args are given.""" workflow = _build_multi_superstep_workflow() diff --git a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py index d9cb44fe01c..fc79a1ae07f 100644 --- a/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py +++ b/python/packages/ag-ui/tests/ag_ui/test_workflow_run.py @@ -98,6 +98,229 @@ def _interrupt_metadata_value(interrupt: dict[str, Any]) -> dict[str, Any]: return cast(dict[str, Any], value) +def test_attach_checkpoint_id_to_interrupts_setdefault() -> None: + """Issue #8150: attach pause checkpoint_id without overwriting an existing one.""" + from agent_framework_ag_ui._workflow_run import _attach_checkpoint_id_to_interrupts + + empty = _attach_checkpoint_id_to_interrupts([{"id": "r1"}], None) + assert empty == [{"id": "r1"}] + + attached = _attach_checkpoint_id_to_interrupts( + [{"id": "r1", "metadata": {"agent_framework": {"type": "workflow_request_info"}}}], + "cp-123", + ) + assert attached[0]["metadata"]["agent_framework"]["checkpoint_id"] == "cp-123" + assert attached[0]["metadata"]["agent_framework"]["type"] == "workflow_request_info" + + preserved = _attach_checkpoint_id_to_interrupts(attached, "cp-other") + assert preserved[0]["metadata"]["agent_framework"]["checkpoint_id"] == "cp-123" + + +@pytest.mark.asyncio +async def test_pause_checkpoint_id_ignores_competing_shared_latest() -> None: + """Prefer this runner's pause checkpoint over a newer shared get_latest() winner.""" + from agent_framework import WorkflowCheckpoint + + from agent_framework_ag_ui._run_common import _build_run_finished_event + from agent_framework_ag_ui._workflow_run import ( + _interrupts_with_pause_checkpoint, + _pause_checkpoint_id_for_interrupts, + ) + + 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"}, + ) + 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, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + first_events = [ + event + async for event in run_workflow_stream( + {"messages": [{"role": "user", "content": "go"}]}, + workflow, + checkpoint_storage=storage, + ) + ] + first_finished = [event for event in first_events if event.type == "RUN_FINISHED"][0] + interrupt_payload = _interrupts_from_run_finished(first_finished) + pause_id = interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] + + # Inject a newer shared checkpoint that does not cover this interrupt (other owner / stale). + competing = WorkflowCheckpoint( + workflow_name=workflow.name, + graph_signature_hash="competing", + pending_request_info_events={}, + timestamp="9999-01-01T00:00:00+00:00", + ) + await storage.save(competing) + latest = await storage.get_latest(workflow_name=workflow.name) + assert latest is not None + assert latest.checkpoint_id == competing.checkpoint_id + + # Workflow-owned baseline from the pause run should already allow advertising pause_id. + resolved = await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=storage, + interrupts=interrupt_payload, + ) + assert resolved == pause_id + + rebuilt = _build_run_finished_event( + "run-1", + "thread-1", + interrupts=await _interrupts_with_pause_checkpoint( + interrupts=interrupt_payload, + workflow=workflow, + checkpoint_storage=storage, + ), + ) + rebuilt_interrupts = _interrupts_from_run_finished(rebuilt) + assert rebuilt_interrupts[0]["metadata"]["agent_framework"]["checkpoint_id"] == pause_id + + +@pytest.mark.asyncio +async def test_builder_checkpoint_storage_attaches_id_without_run_arg() -> None: + """WorkflowBuilder(checkpoint_storage=...) alone must still advertise pause checkpoint_id.""" + + 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"}, + ) + 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, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + # Deliberately omit checkpoint_storage= on the AG-UI entrypoint (builder path only). + events = [ + event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow) + ] + finished = [event for event in events if event.type == "RUN_FINISHED"][0] + interrupt_payload = _interrupts_from_run_finished(finished) + checkpoints = await storage.list_checkpoints(workflow_name=workflow.name) + assert checkpoints + assert interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] == checkpoints[-1].checkpoint_id + + +@pytest.mark.asyncio +async def test_builder_checkpoint_storage_resume_round_trips_without_agui_storage() -> None: + """Pause IDs from builder storage must resume with resume payload even without AG-UI storage.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + await ctx.request_info("need-input", str, request_id="req-1") + + @response_handler + async def handle(self, original_request: str, response: str, ctx: WorkflowContext) -> None: + del original_request + await ctx.yield_output(f"got:{response}") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=ApprovalExecutor(), checkpoint_storage=storage).build() + + pause_events = [ + event async for event in run_workflow_stream({"messages": [{"role": "user", "content": "go"}]}, workflow) + ] + finished = [event for event in pause_events if event.type == "RUN_FINISHED"][0] + pause_id = _interrupts_from_run_finished(finished)[0]["metadata"]["agent_framework"]["checkpoint_id"] + + # Cold resume: omit AG-UI checkpoint_storage; rely on builder storage only. + resume_events = [ + event + async for event in run_workflow_stream( + { + "messages": [], + "resume": {"interrupts": [{"id": "req-1", "value": "ok"}]}, + "forwarded_props": {"checkpoint_id": pause_id}, + }, + workflow, + ) + ] + assert "RUN_ERROR" not in [event.type for event in resume_events] + assert "RUN_FINISHED" in [event.type for event in resume_events] + text = "".join(getattr(event, "delta", "") for event in resume_events if event.type == "TEXT_MESSAGE_CONTENT") + assert "got:ok" in text + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_is_run_scoped_without_storage() -> None: + """Stale runner ids must not be advertised when baseline shows this run did not persist.""" + + class ApprovalExecutor(Executor): + def __init__(self) -> None: + super().__init__(id="approval_executor") + + @handler + async def start(self, message: Any, ctx: WorkflowContext) -> None: + del message + await ctx.request_info("need-input", str, request_id="req-1") + + @response_handler + async def handle(self, original_request: str, response: str, ctx: WorkflowContext) -> None: + del original_request, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + workflow = WorkflowBuilder(start_executor=ApprovalExecutor()).build() + # Simulate a leftover id from a prior run on a workflow with no checkpoint storage. + workflow._runner._previous_checkpoint_id = "stale-from-prior-run" # pyright: ignore[reportPrivateUsage] + + from agent_framework_ag_ui._workflow_run import _pause_checkpoint_id_for_interrupts + + interrupts = [{"id": "req-1", "value": "need-input"}] + workflow._run_baseline_checkpoint_id = "stale-from-prior-run" # pyright: ignore[reportPrivateUsage] + assert ( + await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=None, + interrupts=interrupts, + ) + is None + ) + assert ( + await _pause_checkpoint_id_for_interrupts( + workflow=workflow, + checkpoint_storage=None, + interrupts=interrupts, + baseline_checkpoint_id=None, + ) + == "stale-from-prior-run" + ) + + async def test_workflow_run_maps_custom_and_text_events(): """Custom workflow events and yielded text are mapped to AG-UI events.""" @@ -843,6 +1066,8 @@ async def handle_approval(self, original_request: Content, response: Content, ct ) assert checkpoints, "expected the interrupted run to create a checkpoint" resume_checkpoint_id = checkpoints[-1].checkpoint_id + # Issue #8150: interrupt metadata must carry the pause checkpoint for multi-worker resume. + assert interrupt_payload[0]["metadata"]["agent_framework"]["checkpoint_id"] == resume_checkpoint_id # Resume on a FRESH workflow instance so no pending requests exist in memory until # the checkpoint is restored -- a cold restore, as after a process restart. diff --git a/python/packages/core/AGENTS.md b/python/packages/core/AGENTS.md index d63912a5477..d37fbb1859d 100644 --- a/python/packages/core/AGENTS.md +++ b/python/packages/core/AGENTS.md @@ -73,6 +73,9 @@ agent_framework/ - **`FunctionTool`** - Wraps Python functions as tools with JSON schema generation - **`@tool`** decorator - Converts functions to tools - **`use_function_invocation()`** - Decorator to add automatic function calling to chat clients +- **`_format_tool_parameters`** - Private structured parameter formatter shared by Hyperlight and Monty descriptions. + Returns the effective compact/JSON format and detached parameter data, falling back to full JSON Schema when + compact data cannot preserve constraints. Does not change `FunctionTool.parameters()` or render runtime-specific text. ### Vector stores @@ -157,6 +160,7 @@ The vector store API is experimental under the shared `VECTOR_STORES` feature ID ### Skills (`_skills.py`) +- **Skill frontmatter parsing** - Local and MCP archive skill loaders use PyYAML's `SafeLoader` node composition, not dictionary construction, so duplicate mapping entries remain available for validation and no YAML object constructors run. Recognized root names must be lowercase and unique after YAML decoding. Scalar fields remain text (including numeric/boolean scalar spellings); null root values remain unset. Decoded keys and retained scalar values cannot contain Unicode surrogate code points, which cannot be encoded as UTF-8. Invalid optional metadata mappings or entries warn and are skipped; case-sensitive duplicate metadata keys keep the first valid value. Invalid YAML syntax or invalid recognized root fields reject the skill. Quoting, escapes, multiline folding/chomping, and line-ending normalization follow PyYAML. YAML merge keys (`<<`) are not expanded: they are invalid at the root and skipped with warnings in metadata. - **`Skill`** - Abstract base for a skill definition bundling instructions (`content`) with frontmatter metadata, resources, and scripts. Concrete subclasses (`InlineSkill`, `FileSkill`, `ClassSkill`) accept a `frontmatter=SkillFrontmatter(...)` argument carrying the spec fields. Adding new spec fields is done in one place — on `SkillFrontmatter` — keeping the subclass constructors stable. - **`SkillFrontmatter`** - L1 discovery metadata for a skill (`name`, `description`, `license`, `compatibility`, `allowed_tools`, `metadata`). All fields are mutable plain attributes; the constructor validates `name`, `description`, and `compatibility` against the spec but post-construction assignments are not re-validated. Spec fields are reachable on every skill via `skill.frontmatter`. - **`SkillResource`** - Named supplementary content attached to a skill; holds either static `content` or a dynamic `function` (sync or async). Exactly one must be provided. diff --git a/python/packages/core/agent_framework/_agents.py b/python/packages/core/agent_framework/_agents.py index c4f5da1c0de..205d44b35c6 100644 --- a/python/packages/core/agent_framework/_agents.py +++ b/python/packages/core/agent_framework/_agents.py @@ -1405,8 +1405,10 @@ async def _prepare_run_context( self.context_providers.append(InMemoryHistoryProvider()) active_session = session + framework_created_session = False if active_session is None and self.context_providers: active_session = AgentSession() + framework_created_session = True per_service_call_history_providers = self._resolve_per_service_call_history_providers( session=active_session, @@ -1529,8 +1531,12 @@ async def _prepare_run_context( session_messages: list[Message] = session_context.get_messages(include_input=True) effective_client_kwargs = dict(client_kwargs) if client_kwargs is not None else {} + from ._tools import _APPROVAL_SESSION_IS_AUTHORITATIVE_KEY # pyright: ignore[reportPrivateUsage] + + effective_client_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, None) if active_session is not None: effective_client_kwargs["session"] = active_session + effective_client_kwargs[_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY] = not framework_created_session per_service_call_history_middleware: PerServiceCallHistoryPersistingMiddleware | None = None if per_service_call_history_providers and active_session is not None: per_service_call_history_middleware = PerServiceCallHistoryPersistingMiddleware( diff --git a/python/packages/core/agent_framework/_compaction.py b/python/packages/core/agent_framework/_compaction.py index 227417bb57a..45205866561 100644 --- a/python/packages/core/agent_framework/_compaction.py +++ b/python/packages/core/agent_framework/_compaction.py @@ -30,6 +30,12 @@ GROUP_INDEX_KEY = "index" GROUP_HAS_REASONING_KEY = "has_reasoning" GROUP_TOKEN_COUNT_KEY = "token_count" # ruff:ignore[hardcoded-password-string] # nosec B105 - compaction metadata key, not a credential +GROUP_TOKEN_COUNT_BASIS_KEY = "token_count_basis" # ruff:ignore[hardcoded-password-string] # nosec B105 - compaction metadata key, not a credential +# Bumped whenever the token-estimation serialization changes. Cached counts +# stamped with an older basis are recomputed instead of reused: annotations +# survive Message.to_dict()/from_dict() round-trips, so without the basis a +# transcript annotated before a serialization change keeps its stale counts. +TOKEN_COUNT_BASIS_VERSION = 2 EXCLUDED_KEY = "_excluded" EXCLUDE_REASON_KEY = "_exclude_reason" SUMMARY_OF_MESSAGE_IDS_KEY = "_summary_of_message_ids" @@ -496,7 +502,13 @@ def _write_group_annotation( token_count: int | None = None if existing_raw_annotation is not None: raw_token_count = existing_raw_annotation.get(GROUP_TOKEN_COUNT_KEY) - if isinstance(raw_token_count, int) or raw_token_count is None: + raw_token_count_basis = existing_raw_annotation.get(GROUP_TOKEN_COUNT_BASIS_KEY) + # Preserve a cached count only when it was computed with the current + # serialization basis; counts from older bases are stale and get + # recomputed (see TOKEN_COUNT_BASIS_VERSION). + if raw_token_count_basis == TOKEN_COUNT_BASIS_VERSION and ( + isinstance(raw_token_count, int) or raw_token_count is None + ): token_count = raw_token_count unknown_fields = { key: value @@ -508,6 +520,7 @@ def _write_group_annotation( GROUP_INDEX_KEY, GROUP_HAS_REASONING_KEY, GROUP_TOKEN_COUNT_KEY, + GROUP_TOKEN_COUNT_BASIS_KEY, } } @@ -517,6 +530,7 @@ def _write_group_annotation( GROUP_INDEX_KEY: index, GROUP_HAS_REASONING_KEY: has_reasoning, GROUP_TOKEN_COUNT_KEY: token_count, + GROUP_TOKEN_COUNT_BASIS_KEY: TOKEN_COUNT_BASIS_VERSION if token_count is not None else None, } annotation.update(unknown_fields) message.additional_properties[GROUP_ANNOTATION_KEY] = annotation @@ -550,6 +564,10 @@ def _token_count(message: Message) -> int | None: if annotation is None: return None token_count = annotation.get(GROUP_TOKEN_COUNT_KEY) + # Counts stamped with an older serialization basis are stale: they were + # computed from a different serialized form and must be recomputed. + if annotation.get(GROUP_TOKEN_COUNT_BASIS_KEY) != TOKEN_COUNT_BASIS_VERSION: + return None return token_count if isinstance(token_count, int) else None @@ -558,6 +576,7 @@ def _write_token_count(message: Message, token_count: int) -> None: if annotation is None: return annotation[GROUP_TOKEN_COUNT_KEY] = token_count + annotation[GROUP_TOKEN_COUNT_BASIS_KEY] = TOKEN_COUNT_BASIS_VERSION message.additional_properties[GROUP_ANNOTATION_KEY] = annotation @@ -732,12 +751,68 @@ def annotate_message_groups( return _ordered_group_ids_from_annotations(messages) +_OPAQUE_REASONING_KEYS: Final[frozenset[str]] = frozenset({"encrypted_content"}) + + +def _strip_opaque_reasoning_payload(value: Any) -> Any: + """Recursively drop opaque reasoning members, keeping clear-text payloads.""" + if isinstance(value, dict): + entries = cast("dict[Any, Any]", value) + filtered: dict[Any, Any] = {} + for key, item in entries.items(): + if key in _OPAQUE_REASONING_KEYS: + continue + stripped = _strip_opaque_reasoning_payload(item) + if stripped is not None: + filtered[key] = stripped + return filtered or None + if isinstance(value, list): + entries = cast("list[Any]", value) + kept = [ + stripped + for stripped in (_strip_opaque_reasoning_payload(entry) for entry in entries) + if stripped is not None + ] + return kept or None + return value + + def _serialize_content(content: Content) -> dict[str, Any]: payload = content.to_dict(exclude_none=True) payload.pop("raw_representation", None) # ``items`` mirrors ``result`` for function_result content; exclude it # to avoid double-counting tokens during estimation. payload.pop("items", None) + # ``protected_data`` carries provider reasoning payloads. JSON-serialised + # reasoning_details (Chat Completions) are replayed to the provider as + # clear text -- ``summary``, ``reasoning_text`` and nested ``reasoning.text`` + # are part of the context the provider receives, so only their opaque + # members (``encrypted_content``) are excluded and the clear text stays + # counted. Anything that is not such a JSON structure (Anthropic thinking + # ``signature``, Responses API ``encrypted_content`` blobs) is replayed + # opaquely and never tokenised; exclude it so estimation measures the text + # the model actually sees. + protected_data = payload.get("protected_data") + if isinstance(protected_data, str) and protected_data: + try: + reasoning_payload = json.loads(protected_data) + except ValueError: + reasoning_payload = None + filtered = ( + _strip_opaque_reasoning_payload(reasoning_payload) if isinstance(reasoning_payload, (dict, list)) else None + ) + if filtered is None: + payload.pop("protected_data", None) + else: + payload["protected_data"] = json.dumps(filtered, ensure_ascii=False) + else: + payload.pop("protected_data", None) + additional_properties = payload.get("additional_properties") + if isinstance(additional_properties, dict) and "encrypted_content" in additional_properties: + typed_properties = cast("dict[str, Any]", additional_properties) + payload["additional_properties"] = { + key: value for key, value in typed_properties.items() if key != "encrypted_content" + } return payload diff --git a/python/packages/core/agent_framework/_harness/_mode.py b/python/packages/core/agent_framework/_harness/_mode.py index 11cc018b178..9d8e1650286 100644 --- a/python/packages/core/agent_framework/_harness/_mode.py +++ b/python/packages/core/agent_framework/_harness/_mode.py @@ -8,18 +8,26 @@ from .._sessions import AgentSession, ContextProvider, SessionContext from .._telemetry import FeatureIndex, mark_feature_used -from .._tools import tool +from .._tools import FunctionTool, tool from .._types import Message DEFAULT_MODE_SOURCE_ID = "agent_mode" +_MODE_GET_INSTRUCTIONS = "Use the mode_get tool to check your current operating mode.\n" +_MODE_SET_INSTRUCTIONS = ( + "Use the mode_set tool to switch between modes as your work progresses. " + "Only use mode_set if the user explicitly instructs/allows you to change modes.\n\n" +) +_PLAN_MODE_TRANSITION = ( + "7. When approval is granted, always switch to execute mode (using the `mode_set` tool), " + "and follow the steps for *Execute mode*." +) DEFAULT_MODE_INSTRUCTIONS = ( "## Agent Mode\n\n" "- You can operate in different modes. Depending on the mode you are in, " "you will be required to follow different processes.\n\n" - "Use the mode_get tool to check your current operating mode.\n" - "Use the mode_set tool to switch between modes as your work progresses. " - "Only use mode_set if the user explicitly instructs/allows you to change modes.\n\n" - "You are currently operating in the {current_mode} mode.\n\n" + + _MODE_GET_INSTRUCTIONS + + _MODE_SET_INSTRUCTIONS + + "You are currently operating in the {current_mode} mode.\n\n" "### Mandatory Mode based Workflow\n\n" "For every new substantive user request, including short factual questions, " "your behavior is determined by the mode you are in.\n\n" @@ -49,8 +57,7 @@ "5. Write the plan to a memory file, so that it is retained even if compaction happens. " "Make sure to update the plan file if the user requests changes.\n" "6. Present the plan to the user and ask for approval to switch to execute mode and process the plan.\n" - "7. When approval is granted, always switch to execute mode (using the `mode_set` tool), " - "and follow the steps for *Execute mode*." + + _PLAN_MODE_TRANSITION ), "execute": ( "Determine the type of ask:\n" @@ -157,6 +164,7 @@ def set_agent_mode( *, source_id: str = DEFAULT_MODE_SOURCE_ID, available_modes: Sequence[str] | None = None, + notify: bool = True, ) -> str: """Set the current operating mode in session state. @@ -173,6 +181,9 @@ def set_agent_mode( Keyword Args: source_id: Unique source ID for the provider state. available_modes: Supported modes to validate against. Defaults to the built-in modes. + notify: Whether to notify the agent about the mode change on its next run. Set to ``False`` when the + agent changes mode through a replacement tool and has already observed the tool result. This also + clears any pending external-change notification. Returns: The normalized mode string that was stored. @@ -189,8 +200,11 @@ def set_agent_mode( # prior mode so the next ``before_run`` can inject a user message announcing the switch. Without # that injection, the model often anchors on the earlier ``set_mode`` tool call in the chat history # and keeps behaving as if it were still in that mode — system instructions alone are insufficient. - if isinstance(previous_mode, str) and previous_mode != normalized_mode: - provider_state[_PREVIOUS_MODE_STATE_KEY] = previous_mode + if notify: + if isinstance(previous_mode, str) and previous_mode != normalized_mode: + provider_state[_PREVIOUS_MODE_STATE_KEY] = previous_mode + else: + provider_state.pop(_PREVIOUS_MODE_STATE_KEY, None) return normalized_mode @@ -204,10 +218,13 @@ class AgentModeProvider(ContextProvider): The set of available modes is configurable with ``mode_instructions``. By default, two modes are provided: ``"plan"`` (interactive planning) and ``"execute"`` (autonomous execution). - This provider exposes the following tools to the agent: + By default, this provider exposes the following tools to the agent: - ``mode_set``: Switch the agent's operating mode. - ``mode_get``: Retrieve the agent's current operating mode. + Set ``expose_mode_set`` or ``expose_mode_get`` to ``False`` to omit that tool while retaining mode state + and workflow instructions. Replacement tools can be supplied through the agent's ``tools`` argument. + Public helper functions ``get_agent_mode`` and ``set_agent_mode`` allow external code to programmatically read and change the mode. """ @@ -219,6 +236,8 @@ def __init__( default_mode: str | None = None, mode_instructions: Mapping[str, str] | None = None, instructions: str | None = None, + expose_mode_set: bool = True, + expose_mode_get: bool = True, ) -> None: """Initialize a new agent mode provider. @@ -229,15 +248,30 @@ def __init__( default_mode: Initial mode used when no mode is stored yet. When omitted, the first entry of ``mode_instructions`` is used. mode_instructions: Mapping of supported modes to instructions on when and how to use each mode. + Custom text is not rewritten when tools are hidden. instructions: Custom instructions for using the mode tools. The instructions can contain an ``{available_modes}`` placeholder for the configured list of modes and a ``{current_mode}`` placeholder for the currently active mode. When omitted, the provider uses a default set of instructions. + Default guidance reflects tool exposure; custom text is not rewritten when tools are hidden. + expose_mode_set: Whether to contribute the built-in ``mode_set`` tool. Defaults to ``True``. + When ``False``, the application controls mode changes, optionally through a replacement tool. + expose_mode_get: Whether to contribute the built-in ``mode_get`` tool. Defaults to ``True``. + The current mode remains available in the default instructions and through ``get_agent_mode``. Raises: ValueError: No modes are configured, or the default mode is not configured. """ super().__init__(source_id) - mode_instructions = dict(DEFAULT_MODE_MAP if mode_instructions is None else mode_instructions) + if mode_instructions is None: + mode_instructions = dict(DEFAULT_MODE_MAP) + if not expose_mode_set: + mode_instructions["plan"] = mode_instructions["plan"].replace( + _PLAN_MODE_TRANSITION, + "7. When approval is granted, use the application's configured mode-change mechanism to " + "transition to execute mode. Follow the steps for *Execute mode* only after the mode has changed.", + ) + else: + mode_instructions = dict(mode_instructions) self._mode_display_names = _normalize_available_modes(tuple(mode_instructions)) if not self._mode_display_names: raise ValueError("mode_instructions must contain at least one mode.") @@ -247,6 +281,8 @@ def __init__( self.available_modes = tuple(self._mode_display_names) self.default_mode = _resolve_default_mode(default_mode, available_modes=self._mode_display_names) self.instructions = instructions + self.expose_mode_set = expose_mode_set + self.expose_mode_get = expose_mode_get def _build_instructions(self, current_mode: str) -> str: """Build the mode guidance injected for the current session.""" @@ -255,6 +291,15 @@ def _build_instructions(self, current_mode: str) -> str: for mode, mode_instruction in self.mode_instructions.items() ) instructions = self.instructions or DEFAULT_MODE_INSTRUCTIONS + if not self.instructions: + if not self.expose_mode_get: + instructions = instructions.replace(_MODE_GET_INSTRUCTIONS, "") + if not self.expose_mode_set: + instructions = instructions.replace( + _MODE_SET_INSTRUCTIONS, + "Mode changes are controlled by the application. Use its configured mode-change mechanism " + "only when the user explicitly instructs/allows a mode change.\n\n", + ) return instructions.replace("{available_modes}", mode_lines).replace("{current_mode}", current_mode) async def before_run( @@ -289,11 +334,13 @@ async def before_run( @tool(name="mode_set", approval_mode="never_require") def mode_set(mode: str) -> str: """Switch the agent's operating mode.""" - # The agent invoked the tool itself, so it knows the mode just changed — bypass - # ``set_agent_mode`` to avoid triggering a notification message on the next turn. - normalized_mode = _normalize_mode(mode, available_modes=self._mode_display_names) - tool_state = _get_mode_state(session, source_id=self.source_id) - tool_state["current_mode"] = normalized_mode + normalized_mode = set_agent_mode( + session, + mode, + source_id=self.source_id, + available_modes=self.available_modes, + notify=False, + ) return json.dumps({"mode": normalized_mode, "message": f"Mode changed to '{normalized_mode}'."}) @tool(name="mode_get", approval_mode="never_require") @@ -311,7 +358,12 @@ def mode_get() -> str: self.source_id, [self._build_instructions(current_mode)], ) - context.extend_tools(self.source_id, [mode_set, mode_get]) + tools: list[FunctionTool] = [] + if self.expose_mode_set: + tools.append(mode_set) + if self.expose_mode_get: + tools.append(mode_get) + context.extend_tools(self.source_id, tools) if isinstance(previous_mode, str) and previous_mode != current_mode: # Inject a user-role message announcing the external mode change. System instructions # always render first in the chat history, so the agent can otherwise stay anchored to diff --git a/python/packages/core/agent_framework/_harness/_tool_approval.py b/python/packages/core/agent_framework/_harness/_tool_approval.py index ad81d06451d..4e2db951c77 100644 --- a/python/packages/core/agent_framework/_harness/_tool_approval.py +++ b/python/packages/core/agent_framework/_harness/_tool_approval.py @@ -458,14 +458,22 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: approval_requests: list[Content] = [] buffered_approval_updates: list[AgentResponseUpdate] = [] + streamed_user_input_contents: list[Content] = [] + buffering = False async for update in context.result: update_approval_requests = [ content for content in update.contents if content.type == "function_approval_request" ] - if not update_approval_requests: + if not buffering and not update_approval_requests: + streamed_user_input_contents.extend( + content for content in update.contents if content.user_input_request + ) + if update_approval_requests: + buffering = True + approval_requests.extend(update_approval_requests) + if not buffering: yield update continue - approval_requests.extend(update_approval_requests) buffered_update = copy.copy(update) buffered_update.contents = list(update.contents) buffered_approval_updates.append(buffered_update) @@ -476,7 +484,10 @@ async def _stream() -> AsyncIterable[AgentResponseUpdate]: response_messages = [ Message( role="assistant", - contents=[content for update in buffered_approval_updates for content in update.contents], + contents=[ + *streamed_user_input_contents, + *(content for update in buffered_approval_updates for content in update.contents), + ], ) ] has_other_user_input = self._has_non_approval_user_input(response_messages) diff --git a/python/packages/core/agent_framework/_skills.py b/python/packages/core/agent_framework/_skills.py index b02ec2cf5e5..eb29d6ab5b4 100644 --- a/python/packages/core/agent_framework/_skills.py +++ b/python/packages/core/agent_framework/_skills.py @@ -64,6 +64,9 @@ from pathlib import Path, PurePosixPath from typing import IO, TYPE_CHECKING, Any, ClassVar, Final, Protocol, TypeAlias, TypeVar, cast, runtime_checkable +import yaml +from yaml.nodes import MappingNode, Node, ScalarNode + from ._feature_stage import ExperimentalFeature, experimental from ._filesystem import _is_link_or_reparse_point # pyright: ignore[reportPrivateUsage] from ._sessions import ContextProvider @@ -701,7 +704,9 @@ class SkillFrontmatter: compatibility: Optional compatibility information (≤500 characters). allowed_tools: Optional space-delimited pre-approved tool names. metadata: Optional arbitrary key-value pairs (shallow-copied on - construction to avoid caller-owned dict aliasing). + construction to avoid caller-owned dict aliasing). Keys are case-sensitive. + When parsed from a SKILL.md file, invalid entries are skipped with warnings, + and exact duplicate keys retain the first valid value with warnings. """ def __init__( @@ -1749,125 +1754,29 @@ def __call__( # Matches YAML frontmatter delimited by "---" lines. # The \uFEFF? prefix allows an optional UTF-8 BOM. FRONTMATTER_RE = re.compile( - r"\A\uFEFF?---\s*$(.+?)^---\s*$", + r"\A\uFEFF?---[ \t]*\r?\n(.*?)^---[ \t]*\r?$", re.MULTILINE | re.DOTALL, ) -# Matches top-level YAML "key: value" lines (unindented). Group 1 = key, -# Group 2 = quoted value, Group 3 = unquoted value. Only matches keys at -# column 0 so that indented children (e.g. under "metadata:") are not -# mistakenly captured as top-level fields. -YAML_KV_RE = re.compile( - r"^([\w-]+)\s*:\s*(?:[\"'](.+?)[\"']|(.+?))\s*$", - re.MULTILINE, -) +FRONTMATTER_FIELD_NAMES = { + "name": "name", + "description": "description", + "license": "license", + "compatibility": "compatibility", + "metadata": "metadata", + "allowed-tools": "allowed-tools", +} -# Matches a YAML "metadata:" block followed by indented key-value pairs. -YAML_METADATA_BLOCK_RE = re.compile( - r"^metadata\s*:\s*$\n((?:[ \t]+\S.*\n?)+)", - re.MULTILINE, -) - -# Matches indented "key: value" lines within a metadata block. -YAML_INDENTED_KV_RE = re.compile( - r"^\s+([\w-]+)\s*:\s*(?:[\"'](.+?)[\"']|(.+?))\s*$", - re.MULTILINE, -) +_YAML_TEXT_SCALAR_TAGS = frozenset(f"tag:yaml.org,2002:{kind}" for kind in ("str", "bool", "int", "float", "timestamp")) # Validates skill names: lowercase letters, numbers, hyphens only; # must not start or end with a hyphen, and must not contain consecutive hyphens. VALID_NAME_RE = re.compile(r"^[a-z0-9]([a-z0-9]*-[a-z0-9])*[a-z0-9]*$") -# Block scalar indicator characters recognised by the lightweight YAML parser. -_BLOCK_SCALAR_INDICATORS = ("|", ">") - - -def _parse_yaml_scalar_value(yaml_content: str, kv_match: re.Match[str]) -> str: - """Resolve the scalar value for an unquoted YAML key-value match. - - If the captured value starts with a YAML block scalar indicator (``|`` or - ``>``), the function reads subsequent indented continuation lines, strips - the common leading indentation, and joins them according to the scalar - style (literal preserves newlines, folded replaces them with spaces). - - Chomping indicators are respected per YAML 1.2 §8.1.1.2: - - * ``-`` (strip) — final line break and trailing empty lines excluded - * ``+`` (keep) — final line break and any trailing empty lines preserved - * default (clip) — final line break preserved, trailing empty lines excluded - - For plain (non-block-scalar) values the captured text is returned as-is. - Note: explicit indentation indicators (e.g. ``|2``) are not supported; - indentation is auto-detected from the common leading whitespace. - """ - value: str = kv_match.group(3) - - if not value or value[0] not in _BLOCK_SCALAR_INDICATORS: - return value - - scalar_style = value[0] - keep_trailing_newline = len(value) > 1 and value[1] == "+" - strip_trailing_newline = len(value) > 1 and value[1] == "-" - - # Find the start of the next line after this key-value match. - next_line_start = yaml_content.find("\n", kv_match.end()) - if next_line_start < 0: - return value - next_line_start += 1 # skip the newline character itself - - # Collect indented continuation lines (or blank lines within the block). - block_lines: list[str] = [] - pos = next_line_start - while pos < len(yaml_content): - line_end = yaml_content.find("\n", pos) - if line_end < 0: - line = yaml_content[pos:] - line_end = len(yaml_content) - else: - line = yaml_content[pos:line_end] - - if not line or line.isspace(): - # Blank / whitespace-only lines are part of the block. - block_lines.append("") - pos = line_end + 1 if line_end < len(yaml_content) else line_end - continue - - if line[0] not in (" ", "\t"): - # Non-indented, non-blank line — end of the block. - break - - block_lines.append(line) - pos = line_end + 1 if line_end < len(yaml_content) else line_end - - # Strip trailing blank lines collected from the block. - while block_lines and block_lines[-1] == "": - block_lines.pop() - if not block_lines: - return "" - - # Determine the common leading indentation across non-empty lines. - # Only space/tab characters count as indentation (matches YAML semantics). - def _indent_width(s: str) -> int: - i = 0 - while i < len(s) and s[i] in (" ", "\t"): - i += 1 - return i - - common_indent = min(_indent_width(line) for line in block_lines if line) - normalized = [line[common_indent:] if line else "" for line in block_lines] - - # Literal preserves newlines; folded joins non-empty lines with spaces. - parsed = "\n".join(normalized) if scalar_style == "|" else " ".join(line for line in normalized if line) - - if keep_trailing_newline: - return parsed + "\n" - if strip_trailing_newline: - return parsed - # Clip (default): literal gets a trailing newline, folded does not. - if scalar_style == "|": - return parsed + "\n" - return parsed +def _contains_surrogate_code_points(value: str) -> bool: + """Detect surrogate escapes that PyYAML accepts but UTF-8 cannot encode.""" + return any("\ud800" <= char <= "\udfff" for char in value) # Default system prompt template for advertising available skills to the model. @@ -3600,6 +3509,58 @@ def _validate_skill_metadata( return None + @staticmethod + def _extract_frontmatter_metadata(node: Node | None, skill_file_path: str) -> dict[str, str] | None: + """Read optional metadata without rejecting a skill for invalid entries.""" + # Missing metadata, an empty declaration, or an explicit YAML null leaves metadata unset. + if node is None or (isinstance(node, ScalarNode) and node.tag == "tag:yaml.org,2002:null"): + return None + + # Metadata must be a standard YAML mapping; otherwise warn and ignore it without rejecting the skill. + if not isinstance(node, MappingNode) or node.tag != "tag:yaml.org,2002:map": + logger.warning("SKILL.md at '%s' has invalid metadata; expected a mapping", skill_file_path) + return None + + metadata: dict[str, str] = {} + + for key_node, value_node in node.value: + # Accept scalar keys as text, but reject collections, nulls, and unsupported YAML types. + if ( + not isinstance(key_node, ScalarNode) + or key_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(key_node.value) + ): + logger.warning("SKILL.md at '%s' has an invalid metadata key; skipping entry", skill_file_path) + continue + + key = key_node.value + # Keep the first valid value for each exact, case-sensitive key; warn about later duplicates. + if key in metadata: + logger.warning( + "SKILL.md at '%s' contains duplicate metadata key '%s'; keeping the first value", + skill_file_path, + key, + ) + continue + + # The Agent Skills spec requires string-to-string metadata, not arrays or nested mappings. + # Keep supported scalars as text; warn and skip invalid values rather than rejecting the skill. + if ( + not isinstance(value_node, ScalarNode) + or value_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(value_node.value) + ): + logger.warning( + "SKILL.md at '%s' has an invalid metadata value for '%s'; expected a text scalar; skipping entry", + skill_file_path, + key, + ) + continue + + metadata[key] = value_node.value + + return metadata + @staticmethod def _extract_frontmatter( content: str, @@ -3610,7 +3571,13 @@ def _extract_frontmatter( Parses the ``---``-delimited frontmatter block for all `agentskills.io specification `_ fields: ``name``, ``description``, ``license``, ``compatibility``, - ``allowed-tools``, and ``metadata``. + ``allowed-tools``, and ``metadata``. Recognized top-level fields must + use lowercase names and must not be repeated, including equivalent quoted + or escaped spellings. YAML scalars are retained as text; null root values + remain unset. Within the optional metadata mapping, keys are case-sensitive; + exact duplicates retain the first valid value and produce warnings. + Invalid metadata entries are skipped with warnings. Invalid YAML syntax + or invalid recognized root fields reject the skill. Args: content: Raw text content of the SKILL.md file. @@ -3625,41 +3592,81 @@ def _extract_frontmatter( logger.error("SKILL.md at '%s' does not contain valid YAML frontmatter delimited by '---'", skill_file_path) return None - yaml_content = match.group(1).strip() - name: str | None = None - description: str | None = None - license_value: str | None = None - compatibility: str | None = None - allowed_tools: str | None = None - - for kv_match in YAML_KV_RE.finditer(yaml_content): - key = kv_match.group(1) - value = ( - kv_match.group(2) if kv_match.group(2) is not None else _parse_yaml_scalar_value(yaml_content, kv_match) - ) + try: + # compose() keeps duplicate keys, unlike safe_load(), which silently keeps only the last value. + # The Pyright ignores are needed because PyYAML's type definitions do not fully describe compose(). + root = yaml.compose(match.group(1), Loader=yaml.SafeLoader) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + except (yaml.YAMLError, RecursionError, ValueError, OverflowError) as exc: + # PyYAML raises ValueError or OverflowError for Unicode escapes outside the valid code-point range. + logger.error("SKILL.md at '%s' contains invalid YAML frontmatter: %s", skill_file_path, exc) + return None + + if not isinstance(root, MappingNode) or root.tag != "tag:yaml.org,2002:map": + logger.error("SKILL.md at '%s' must contain a YAML frontmatter mapping", skill_file_path) + return None + + fields: dict[str, str] = {} + metadata_node: Node | None = None + seen_fields: set[str] = set() + + for key_node, value_node in root.value: + # Accept scalar keys as text, but reject collections, nulls, and unsupported YAML types. + if ( + not isinstance(key_node, ScalarNode) + or key_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(key_node.value) + ): + logger.error("SKILL.md at '%s' has an invalid frontmatter property name", skill_file_path) + return None + + key = key_node.value + canonical_key = FRONTMATTER_FIELD_NAMES.get(key.lower()) + + # Unknown fields are intentionally excluded from this validation for forward compatibility. + if canonical_key is None: + continue + + if key != canonical_key: + logger.error( + "SKILL.md at '%s' uses incorrectly cased frontmatter field '%s'; expected '%s'", + skill_file_path, + key, + canonical_key, + ) + return None + + if key in seen_fields: + logger.error("SKILL.md at '%s' contains duplicate frontmatter field '%s'", skill_file_path, key) + return None + + seen_fields.add(key) + + if key == "metadata": + metadata_node = value_node + continue + + # Leave empty declarations and explicit YAML nulls unset; required fields are validated below. + if isinstance(value_node, ScalarNode) and value_node.tag == "tag:yaml.org,2002:null": + continue + + # Reject lists, mappings, and unsupported YAML types because these fields are stored as text. + if ( + not isinstance(value_node, ScalarNode) + or value_node.tag not in _YAML_TEXT_SCALAR_TAGS + or _contains_surrogate_code_points(value_node.value) + ): + logger.error( + "SKILL.md at '%s' has an invalid '%s' value; expected a text scalar", + skill_file_path, + key, + ) + return None - key_lower = key.lower() - if key_lower == "name": - name = value - elif key_lower == "description": - description = value - elif key_lower == "license": - license_value = value - elif key_lower == "compatibility": - compatibility = value - elif key_lower == "allowed-tools": - allowed_tools = value - - # Parse metadata block (indented key-value pairs under "metadata:"). - metadata: dict[str, str] | None = None - metadata_match = YAML_METADATA_BLOCK_RE.search(yaml_content) - if metadata_match: - metadata = {} - for kv_match in YAML_INDENTED_KV_RE.finditer(metadata_match.group(1)): - mk = kv_match.group(1) - mv = kv_match.group(2) if kv_match.group(2) is not None else kv_match.group(3) - metadata[mk] = mv + fields[key] = value_node.value + name = fields.get("name") + description = fields.get("description") + compatibility = fields.get("compatibility") error = FileSkillsSource._validate_skill_metadata(name, description, skill_file_path, compatibility) if error: logger.error(error) @@ -3670,10 +3677,10 @@ def _extract_frontmatter( return SkillFrontmatter( name=cast(str, name), description=cast(str, description), - license=license_value, + license=fields.get("license"), compatibility=compatibility, - allowed_tools=allowed_tools, - metadata=metadata, + allowed_tools=fields.get("allowed-tools"), + metadata=FileSkillsSource._extract_frontmatter_metadata(metadata_node, skill_file_path), ) @staticmethod diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 1caa4d82565..2261768cdd1 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/python/packages/core/agent_framework/_tools.py @@ -105,15 +105,12 @@ def _generate_function_call_occurrence_id() -> str: DEFAULT_MAX_CONSECUTIVE_ERRORS_PER_REQUEST: Final[int] = 3 SHELL_TOOL_KIND_VALUE: Final[str] = "shell" _TOOL_APPROVAL_STATE_KEY: Final[str] = "tool_approval" -_RUN_LOCAL_MIDDLEWARE_SESSION_ATTR: Final[str] = "_run_local_function_middleware_session" +_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY: Final[str] = "_approval_session_is_authoritative" def _has_authoritative_approval_session(invocation_session: AgentSession | None) -> bool: """Return whether approval state belongs to a caller-owned session.""" - return ( - invocation_session is not None - and getattr(invocation_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, False) is not True - ) + return invocation_session is not None _ALREADY_APPROVED_APPROVAL_REQUEST_GROUPS_KEY: Final[str] = "already_approved_approval_request_groups" @@ -403,6 +400,62 @@ def _annotation_includes_function_invocation_context(annotation: Any) -> bool: ) +def _format_tool_parameters( # pyright: ignore[reportUnusedFunction] + parameters: dict[str, Any], + *, + parameter_format: Literal["compact", "json"], +) -> tuple[Literal["compact", "json"], dict[str, Any]]: + """Return the effective format and detached parameter data for tool descriptions. + + Compact data maps parameter names to scalar type, requiredness, and optional + description, enum, and default metadata. Schemas with unrepresented constraints + retain their full JSON Schema so callers can explain the fallback to the model. + """ + if parameter_format not in ("compact", "json"): + raise ValueError("parameter_format must be 'compact' or 'json'.") + + if parameter_format == "json": + return "json", copy.deepcopy(parameters) + + properties = parameters.get("properties") + required = parameters.get("required", []) + if ( + parameters.get("type") != "object" + or parameters.keys() - {"type", "properties", "required", "title", "description"} + or not isinstance(properties, dict) + or not isinstance(required, list) + ): + return "json", copy.deepcopy(parameters) + + property_schemas = cast(dict[object, Any], properties) + required_names = cast(list[object], required) + if not all(isinstance(name, str) and name in property_schemas for name in required_names): + return "json", copy.deepcopy(parameters) + + compact: dict[str, Any] = {} + for name, property_schema in property_schemas.items(): + if not isinstance(name, str) or not isinstance(property_schema, dict): + return "json", copy.deepcopy(parameters) + + schema = cast(dict[str, Any], property_schema) + if schema.get("type") not in ("string", "integer", "number", "boolean", "null") or schema.keys() - { + "type", + "title", + "description", + "enum", + "default", + }: + return "json", copy.deepcopy(parameters) + + compact[name] = { + "type": schema["type"], + "required": name in required_names, + **{key: copy.deepcopy(schema[key]) for key in ("description", "enum", "default") if key in schema}, + } + + return "compact", compact + + ClassT = TypeVar("ClassT", bound="SerializationMixin") @@ -2257,6 +2310,7 @@ async def _try_execute_function_call_groups( tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]], config: FunctionInvocationConfiguration, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, middleware_pipeline: FunctionMiddlewarePipeline | None = None, host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> tuple[list[list[Content]], bool]: @@ -2268,6 +2322,7 @@ async def _try_execute_function_call_groups( tools: The tools available for execution. config: Configuration for function invocation. invocation_session: The agent session for this invocation, if any. + approval_session_is_authoritative: Whether the invocation session owns persistent approval state. middleware_pipeline: Optional middleware pipeline to apply during execution. host_payload_budget: Shared request budget for retained Host-only function result payloads. @@ -2278,6 +2333,8 @@ async def _try_execute_function_call_groups( """ from ._types import Content + approval_session = invocation_session if approval_session_is_authoritative else None + # Normalize the batch to calls owned by this layer before making any control-flow decision. function_calls = [ function_call @@ -2363,18 +2420,18 @@ async def _try_execute_function_call_groups( visible_requests.append(approval_request) pause_groups.append([approval_request]) continue - if not _has_authoritative_approval_session(invocation_session): + if not _has_authoritative_approval_session(approval_session): visible_requests.append(approval_request) pause_groups.append([approval_request]) continue already_approved_requests.append(approval_request) _store_already_approved_approval_requests( - invocation_session, + approval_session, visible_requests, already_approved_requests, ) - _store_pending_approval_requests(invocation_session, visible_requests) - _store_pending_mixed_pause_batch(invocation_session, pause_groups) + _store_pending_approval_requests(approval_session, visible_requests) + _store_pending_mixed_pause_batch(approval_session, pause_groups) return pause_groups, False if has_declaration_only_call: # Declaration-only calls are returned as user input rather than executed locally. @@ -2470,6 +2527,7 @@ async def _execute_function_calls( options: dict[str, Any] | None, config: FunctionInvocationConfiguration, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, middleware_pipeline: FunctionMiddlewarePipeline | None = None, host_payload_budget: _FunctionResultPayloadBudget | None = None, ) -> _FunctionExecutionBatch: @@ -2481,6 +2539,7 @@ async def _execute_function_calls( function_calls=function_calls, tools=tools, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=middleware_pipeline, config=config, host_payload_budget=host_payload_budget, @@ -2966,6 +3025,110 @@ def _same_mixed_pause_response(left: Mapping[str, Any], right: Mapping[str, Any] return left_payload == right_payload +def _match_mixed_pause_responses( + items: list[dict[str, Any]], + responses: Sequence[Content], + *, + approval_response_binder: Callable[[Content], Content | None] | None = None, +) -> tuple[set[int], bool, list[Content], set[int]]: + """Match one complete mixed pause batch without depending on its storage source.""" + approval_items: dict[str, int | None] = {} + host_items_by_occurrence: dict[str, int] = {} + host_items_by_call: dict[str, list[int]] = {} + for index, item in enumerate(items): + request = _content_from_state(item.get("request")) + if request is None: + continue + if item.get("kind") == "approval": + for identity in ( + request.id, + request.function_call.id if request.function_call is not None else None, + ): + if identity is None: + continue + if identity not in approval_items: + approval_items[identity] = index + elif approval_items[identity] != index: + approval_items[identity] = None + elif item.get("kind") == "host" and request.call_id is not None: + host_items_by_call.setdefault(request.call_id, []).append(index) + if request.id is not None: + host_items_by_occurrence[request.id] = index + + matched_content_ids: set[int] = set() + for response in responses: + item_index: int | None = None + if response.type == "function_approval_response": + candidate = approval_response_binder(response) if approval_response_binder is not None else response + if candidate is None: + continue + response_identities = { + str(identity) + for identity in ( + candidate.additional_properties.get(_APPROVAL_REQUEST_ID_KEY), + candidate.id, + ) + if identity is not None + } + matching_indexes = { + matched_index + for identity in response_identities + if (matched_index := approval_items.get(identity)) is not None + } + if len(matching_indexes) == 1: + item_index = matching_indexes.pop() + elif response.type == "function_result" and response.call_id in host_items_by_call: + candidate = response + if response.id is not None: + item_index = host_items_by_occurrence.get(response.id) + request = _content_from_state(items[item_index].get("request")) if item_index is not None else None + if request is None or request.call_id != response.call_id: + item_index = None + else: + matching_indexes = [ + pending_index + for pending_index in host_items_by_call[response.call_id] + if items[pending_index].get("response") is None + or ( + isinstance(items[pending_index].get("response"), Mapping) + and _same_mixed_pause_response( + cast(Mapping[str, Any], items[pending_index]["response"]), + response.to_dict(), + ) + ) + ] + if len(matching_indexes) == 1: + item_index = matching_indexes[0] + else: + continue + + if item_index is None: + continue + candidate_state = candidate.to_dict() + stored_response = items[item_index].get("response") + if stored_response is not None and ( + not isinstance(stored_response, Mapping) + or not _same_mixed_pause_response(cast(Mapping[str, Any], stored_response), candidate_state) + ): + raise RuntimeError(f"Conflicting response for mixed pause occurrence {candidate.id!r}.") + items[item_index]["response"] = candidate_state + matched_content_ids.add(id(response)) + + if any(item.get("response") is None for item in items): + return matched_content_ids, True, [], set() + + ordered_responses: list[Content] = [] + host_result_ids: set[int] = set() + for item in items: + response = _content_from_state(item.get("response")) + if response is None: + return matched_content_ids, True, [], set() + ordered_responses.append(response) + if item.get("kind") == "host": + host_result_ids.add(id(response)) + return matched_content_ids, False, ordered_responses, host_result_ids + + def _stage_pending_mixed_pause_responses( messages: list[Message], invocation_session: AgentSession | None, @@ -2989,78 +3152,24 @@ def _stage_pending_mixed_pause_responses( state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) return False, False, set() - approval_items: dict[str, dict[str, Any]] = {} - host_items_by_occurrence: dict[str, dict[str, Any]] = {} - host_items_by_call: dict[str, list[dict[str, Any]]] = {} - for item in items: - request = _content_from_state(item.get("request")) - if request is None: - continue - if item.get("kind") == "approval": - for identity in ( - request.id, - request.function_call.id if request.function_call is not None else None, - ): - if identity is not None: - approval_items[identity] = item - elif item.get("kind") == "host" and request.call_id is not None: - host_items_by_call.setdefault(request.call_id, []).append(item) - if request.id is not None: - host_items_by_occurrence[request.id] = item - - matched_content_ids: set[int] = set() - for message in messages: - for content in message.contents: - item: dict[str, Any] | None = None - if content.type == "function_approval_response": - rebound = _bind_approval_response_to_pending_request( - content, - invocation_session, - consume=False, - ) - if rebound is None: - continue - request_id = rebound.additional_properties.get(_APPROVAL_REQUEST_ID_KEY) - item = approval_items.get(str(request_id)) if request_id is not None else None - if item is None and rebound.id is not None: - item = approval_items.get(rebound.id) - candidate = rebound - elif content.type == "function_result" and content.call_id in host_items_by_call: - candidate = content - if content.id is not None: - item = host_items_by_occurrence.get(content.id) - request = _content_from_state(item.get("request")) if item is not None else None - if request is None or request.call_id != content.call_id: - item = None - else: - matching_items = [ - pending_item - for pending_item in host_items_by_call[content.call_id] - if pending_item.get("response") is None - or ( - isinstance(pending_item.get("response"), Mapping) - and _same_mixed_pause_response( - cast(Mapping[str, Any], pending_item["response"]), - content.to_dict(), - ) - ) - ] - if len(matching_items) == 1: - item = matching_items[0] - else: - continue + def bind_approval_response(response: Content) -> Content | None: + return _bind_approval_response_to_pending_request( + response, + invocation_session, + consume=False, + ) - if item is None: - continue - candidate_state = candidate.to_dict() - stored_response = item.get("response") - if stored_response is not None and ( - not isinstance(stored_response, Mapping) - or not _same_mixed_pause_response(cast(Mapping[str, Any], stored_response), candidate_state) - ): - raise RuntimeError(f"Conflicting response for mixed pause occurrence {candidate.id!r}.") - item["response"] = candidate_state - matched_content_ids.add(id(content)) + responses = [ + content + for message in messages + for content in message.contents + if content.type in {"function_approval_response", "function_result"} + ] + matched_content_ids, incomplete, ordered_responses, host_result_ids = _match_mixed_pause_responses( + items, + responses, + approval_response_binder=bind_approval_response, + ) if matched_content_ids: filtered_messages: list[Message] = [] @@ -3071,18 +3180,9 @@ def _stage_pending_mixed_pause_responses( messages[:] = filtered_messages state[_PENDING_MIXED_PAUSE_BATCH_KEY] = {"items": items} - if any(item.get("response") is None for item in items): + if incomplete: return True, False, set() - ordered_responses: list[Content] = [] - host_result_ids: set[int] = set() - for item in items: - response = _content_from_state(item.get("response")) - if response is None: - return True, False, set() - ordered_responses.append(response) - if item.get("kind") == "host": - host_result_ids.add(id(response)) messages.append(Message(role="user", contents=ordered_responses)) return False, True, host_result_ids @@ -3090,145 +3190,99 @@ def _stage_pending_mixed_pause_responses( def _stateless_mixed_pause_batch_status( messages: list[Message], ) -> tuple[bool, set[int]]: - """Validate only the latest stateless mixed batch and order its responses.""" + """Validate the latest unresolved stateless mixed batch and order its responses.""" from ._types import Message - approval_requests: list[Content] = [] - host_requests: list[Content] = [] - approval_responses: list[Content] = [] - host_responses: list[Content] = [] - pause_order_contents: list[Content] = [] - batch_index: int | None = None - for index in range(len(messages) - 1, -1, -1): - message = messages[index] - if any(content.type == "function_approval_request" for content in message.contents) and any( - content.type == "function_call" and content.user_input_request for content in message.contents - ): - batch_index = index - break - if batch_index is None: - return False, set() - - for content in messages[batch_index].contents: - if content.type == "function_call": - pause_order_contents.append(content) - if content.user_input_request: - host_requests.append(content) - elif content.type == "function_approval_request": - pause_order_contents.append(content) - approval_requests.append(content) - for message in messages[batch_index + 1 :]: - for content in message.contents: - if content.type == "function_approval_response": - approval_responses.append(content) - elif content.type == "function_result": - host_responses.append(content) + flattened_contents = [content for message in messages for content in message.contents] + request_batches: list[tuple[int, list[dict[str, Any]], set[str]]] = [] + batch_items: list[dict[str, Any]] = [] + batch_kinds: set[str] = set() + last_request_index = -1 - if not approval_requests or not host_requests or not (approval_responses or host_responses): - return False, set() + def answers_current_batch(content: Content) -> bool: + if content.type not in {"function_approval_response", "function_result"}: + return False + for item in batch_items: + request = _content_from_state(item.get("request")) + if request is None: + continue + if item.get("kind") == "approval": + function_call = request.function_call + if content.type == "function_approval_response" and content.id in { + request.id, + function_call.id if function_call is not None else None, + }: + return True + if ( + content.type == "function_result" + and function_call is not None + and content.call_id == function_call.call_id + ): + return True + elif ( + item.get("kind") == "host" + and content.type == "function_result" + and content.call_id == request.call_id + and (content.id is None or request.id is None or content.id == request.id) + ): + return True + return False - approval_request_identities: list[set[str]] = [] - identity_owners: dict[str, int] = {} - for index, request in enumerate(approval_requests): - identities = { - identity - for identity in ( - request.id, - request.function_call.id if request.function_call is not None else None, - ) - if identity is not None - } - for identity in identities: - if identity in identity_owners and identity_owners[identity] != index: - return True, set() - identity_owners[identity] = index - approval_request_identities.append(identities) - - unmatched_approval_requests = set(range(len(approval_requests))) - matched_approval_responses: dict[int, Content] = {} - matched_approval_response_ids: set[int] = set() - for response in approval_responses: - if response.id is None: + for content_index, content in enumerate(flattened_contents): + if batch_items and answers_current_batch(content): + request_batches.append((last_request_index, batch_items, batch_kinds)) + batch_items = [] + batch_kinds = set() + if content.type == "function_approval_request": + kind = "approval" + elif content.type == "function_call" and content.user_input_request: + kind = "host" + else: continue - matching_indexes = [ - index for index, identities in enumerate(approval_request_identities) if response.id in identities - ] - if len(matching_indexes) != 1: + last_request_index = content_index + batch_kinds.add(kind) + batch_items.append({"kind": kind, "request": content.to_dict()}) + if batch_items: + request_batches.append((last_request_index, batch_items, batch_kinds)) + + for batch_end, items, kinds in reversed(request_batches): + if kinds != {"approval", "host"}: continue - matching_index = matching_indexes[0] - previous_response = matched_approval_responses.get(matching_index) - if previous_response is not None and previous_response.to_dict() != response.to_dict(): - function_call = approval_requests[matching_index].function_call - occurrence_id = function_call.id if function_call is not None else approval_requests[matching_index].id - raise RuntimeError(f"Conflicting approval response for occurrence {occurrence_id!r}.") - matched_approval_responses[matching_index] = response - unmatched_approval_requests.discard(matching_index) - matched_approval_response_ids.add(id(response)) - - unmatched_host_requests = list(host_requests) - matched_host_responses: dict[int, Content] = {} - matched_host_result_ids: set[int] = set() - for response in (candidate for candidate in host_responses if candidate.id is not None): - matching_index = next( - ( - index - for index, request in enumerate(unmatched_host_requests) - if request.id == response.id and request.call_id == response.call_id - ), - None, - ) - if matching_index is not None: - request = unmatched_host_requests.pop(matching_index) - matched_host_responses[id(request)] = response - matched_host_result_ids.add(id(response)) - - for response in (candidate for candidate in host_responses if candidate.id is None): - matching_indexes = [ - index for index, request in enumerate(unmatched_host_requests) if request.call_id == response.call_id + responses = [ + content + for content in flattened_contents[batch_end + 1 :] + if content.type in {"function_approval_response", "function_result"} ] - if len(matching_indexes) == 1: - request = unmatched_host_requests.pop(matching_indexes[0]) - matched_host_responses[id(request)] = response - matched_host_result_ids.add(id(response)) - - if unmatched_approval_requests or unmatched_host_requests: - return True, matched_host_result_ids - - approval_indexes_by_occurrence_id = { - request.function_call.id: index - for index, request in enumerate(approval_requests) - if request.function_call is not None and request.function_call.id is not None - } - approval_indexes_by_request_object = {id(request): index for index, request in enumerate(approval_requests)} - ordered_responses: list[Content] = [] - ordered_approval_indexes: set[int] = set() - for content in pause_order_contents: - host_response = matched_host_responses.get(id(content)) - if host_response is not None: - ordered_responses.append(host_response) + matched_response_ids, incomplete, ordered_responses, host_result_ids = _match_mixed_pause_responses( + items, + responses, + ) + if incomplete: + return True, host_result_ids + + pending_approval_response_ids = { + id(response) + for response in _collect_approval_responses( + messages, + non_approval_result_ids={ + id(content) + for content in responses + if content.type == "function_result" and id(content) in matched_response_ids + }, + ).values() + } + if matched_response_ids.isdisjoint(pending_approval_response_ids): continue - approval_index = approval_indexes_by_request_object.get(id(content)) - if approval_index is None and content.type == "function_call" and content.id is not None: - approval_index = approval_indexes_by_occurrence_id.get(content.id) - if approval_index is not None and approval_index not in ordered_approval_indexes: - ordered_responses.append(matched_approval_responses[approval_index]) - ordered_approval_indexes.add(approval_index) - - ordered_responses.extend( - matched_approval_responses[index] - for index in range(len(approval_requests)) - if index not in ordered_approval_indexes - ) - matched_response_ids = {id(response) for response in ordered_responses} | matched_approval_response_ids - filtered_messages: list[Message] = [] - for message in messages: - message.contents = [content for content in message.contents if id(content) not in matched_response_ids] - if message.contents: - filtered_messages.append(message) - filtered_messages.append(Message(role="user", contents=ordered_responses)) - messages[:] = filtered_messages - return False, matched_host_result_ids + filtered_messages: list[Message] = [] + for message in messages: + message.contents = [content for content in message.contents if id(content) not in matched_response_ids] + if message.contents: + filtered_messages.append(message) + filtered_messages.append(Message(role="user", contents=ordered_responses)) + messages[:] = filtered_messages + return False, host_result_ids + return False, set() def _collect_approval_responses( @@ -4018,6 +4072,7 @@ async def _resolve_approval_responses( max_errors: int, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, middleware_pipeline: FunctionMiddlewarePipeline | None = None, settle_dangling_calls: Callable[[Sequence[Content]], Awaitable[None]] | None = None, ) -> _FunctionProcessingResult: @@ -4031,11 +4086,12 @@ async def _resolve_approval_responses( from ._middleware import MiddlewareFailure from ._types import Message + approval_session = invocation_session if approval_session_is_authoritative else None completed_mixed_batch = False - if _has_authoritative_approval_session(invocation_session): + if _has_authoritative_approval_session(approval_session): incomplete_mixed_batch, completed_mixed_batch, host_result_ids = _stage_pending_mixed_pause_responses( prepared_messages, - invocation_session, + approval_session, ) if incomplete_mixed_batch: return _FunctionProcessingResult(errors_in_a_row=errors_in_a_row, action="return") @@ -4047,13 +4103,13 @@ async def _resolve_approval_responses( ) active_pending_ids = ( - set(_load_pending_approval_requests(invocation_session)) - if _has_authoritative_approval_session(invocation_session) + set(_load_pending_approval_requests(approval_session)) + if _has_authoritative_approval_session(approval_session) else None ) - _bind_approval_responses_to_pending_requests(prepared_messages, invocation_session) + _bind_approval_responses_to_pending_requests(prepared_messages, approval_session) if completed_mixed_batch: - state = _get_tool_approval_state(invocation_session, create=False) + state = _get_tool_approval_state(approval_session, create=False) if state is not None: state.pop(_PENDING_MIXED_PAUSE_BATCH_KEY, None) @@ -4066,7 +4122,7 @@ async def _resolve_approval_responses( } if already_approved_responses := _pop_already_approved_approval_responses( - invocation_session, + approval_session, explicit_approval_response_ids, ): prepared_messages.append(Message(role="user", contents=already_approved_responses)) @@ -4151,7 +4207,7 @@ async def _resolve_approval_responses( terminal_contents.extend( request for request_id, request in pending_by_id.items() if request_id not in surfaced_request_ids ) - _store_pending_approval_requests(invocation_session, list(pending_by_id.values())) + _store_pending_approval_requests(approval_session, list(pending_by_id.values())) # 5. Return role-correct output and tell the outer loop whether to return, stop tools, or call the model. executed_function_count = len(execution_result_groups) @@ -4182,8 +4238,10 @@ async def _process_model_function_calls( max_errors: int, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None = None, + approval_session_is_authoritative: bool = True, ) -> _FunctionProcessingResult: """Execute function calls from a newly completed model response.""" + approval_session = invocation_session if approval_session_is_authoritative else None approval_requests = [ content for message in response.messages @@ -4195,7 +4253,7 @@ async def _process_model_function_calls( function_calls = _extract_function_calls(response) if not (function_calls and tools): if approval_requests: - _store_pending_approval_requests(invocation_session, approval_requests) + _store_pending_approval_requests(approval_session, approval_requests) return _FunctionProcessingResult(errors_in_a_row=errors_in_a_row, action="return") # 2. Execute the batch once while preserving each call's result group. @@ -4224,7 +4282,7 @@ async def _process_model_function_calls( if content.type == "function_approval_request" ] if returned_approval_requests: - _store_pending_approval_requests(invocation_session, returned_approval_requests) + _store_pending_approval_requests(approval_session, returned_approval_requests) return processing_result @@ -4385,6 +4443,7 @@ async def _get_response_with_function_invocation( tokenizer: TokenizerProtocol | None, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None, + approval_session_is_authoritative: bool, budget_state: dict[str, Any], max_errors: int, middleware_pipeline: FunctionMiddlewarePipeline | None = None, @@ -4434,6 +4493,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=middleware_pipeline, settle_dangling_calls=settle_approval_replay_calls, ) @@ -4507,6 +4567,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, ) except MiddlewareFailure: # Fail-closed abort: before propagating, settle the batch's calls on a @@ -4598,6 +4659,7 @@ async def _stream_response_with_function_invocation( tokenizer: TokenizerProtocol | None, execute_function_calls: _FunctionCallExecutor, invocation_session: AgentSession | None, + approval_session_is_authoritative: bool, budget_state: dict[str, Any], max_errors: int, middleware_pipeline: FunctionMiddlewarePipeline | None = None, @@ -4644,6 +4706,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=middleware_pipeline, settle_dangling_calls=settle_approval_replay_calls, ) @@ -4777,6 +4840,7 @@ async def settle_approval_replay_calls(function_calls: Sequence[Content]) -> Non max_errors=max_errors, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, ) except MiddlewareFailure: # See the non-streaming loop: settle a service-managed conversation's @@ -4969,9 +5033,12 @@ def get_response( raw_session = request_kwargs.get("session") invocation_session = raw_session if isinstance(raw_session, _AgentSession) else None + approval_session_is_authoritative = ( + request_kwargs.pop(_APPROVAL_SESSION_IS_AUTHORITATIVE_KEY, True) is not False + ) if invocation_session is None and requires_session_state: invocation_session = _AgentSession() - setattr(invocation_session, _RUN_LOCAL_MIDDLEWARE_SESSION_ATTR, True) + approval_session_is_authoritative = False # Bind one executor with the run's custom arguments, middleware, configuration, and session. execute_function_calls = partial( @@ -4979,6 +5046,7 @@ def get_response( custom_args=additional_function_arguments, config=self.function_invocation_configuration, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, middleware_pipeline=function_middleware_pipeline, host_payload_budget=host_payload_budget, ) @@ -5018,6 +5086,7 @@ def get_response( tokenizer=tokenizer, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, budget_state=budget_state, max_errors=max_errors, middleware_pipeline=function_middleware_pipeline, @@ -5041,6 +5110,7 @@ def finalize_stream(updates: Sequence[ChatResponseUpdate]) -> ChatResponse[Any]: tokenizer=tokenizer, execute_function_calls=execute_function_calls, invocation_session=invocation_session, + approval_session_is_authoritative=approval_session_is_authoritative, budget_state=budget_state, max_errors=max_errors, middleware_pipeline=function_middleware_pipeline, diff --git a/python/packages/core/agent_framework/_workflows/_workflow.py b/python/packages/core/agent_framework/_workflows/_workflow.py index afbf45af11b..a24a2fade38 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow.py +++ b/python/packages/core/agent_framework/_workflows/_workflow.py @@ -410,6 +410,12 @@ def __init__( # so a subsequent ``run()`` is allowed. self._active_run: weakref.ref[ResponseStream[WorkflowEvent, WorkflowRunResult]] | None = None + # Run-scoped pause checkpoint bookkeeping (owned by Workflow, not callers). + # Captured at the start of each ``_run_core`` so ``resolve_pause_checkpoint_id`` + # can tell a newly persisted pause from a leftover / restored id. + self._run_baseline_checkpoint_id: str | None = None + self._restored_checkpoint_id: str | None = None + @property def status(self) -> WorkflowRunState: """Return the current run-level status of this workflow instance. @@ -907,6 +913,12 @@ async def _run_core( if checkpoint_storage is not None: self._runner.context.set_runtime_checkpoint_storage(checkpoint_storage) + # Own the run boundary for pause-checkpoint resolution: capture the runner id + # before this run advances it, and remember an incoming restore so it is not + # re-advertised if a later pause save fails. + self._run_baseline_checkpoint_id = self.get_last_checkpoint_id() + self._restored_checkpoint_id = str(checkpoint_id) if checkpoint_id is not None else None + try: # Async validation: a fresh-message run is only allowed when the # runner context has fully drained from any prior run. If it still @@ -1308,6 +1320,93 @@ def output_types(self) -> list[type[Any] | types.UnionType]: return list(output_types) + def get_last_checkpoint_id(self) -> str | None: + """Return the checkpoint id last persisted or restored by this workflow runner.""" + checkpoint_id = self._runner._previous_checkpoint_id # pyright: ignore[reportPrivateUsage] + return str(checkpoint_id) if checkpoint_id is not None else None + + async def resolve_pause_checkpoint_id( + self, + request_ids: Collection[str], + *, + checkpoint_storage: CheckpointStorage | None = None, + known_checkpoint_id: str | None = None, + baseline_checkpoint_id: str | None | object = _MISSING, + ) -> str | None: + """Resolve the persisted pause checkpoint that covers ``request_ids``. + + Prefers this runner's last-saved id when it advanced past the run baseline + (captured automatically at ``run()`` start), over shared ``get_latest(workflow_name=...)``. + Storage precedence is the run argument, else the runner's effective + (runtime / builder) storage. When storage is available, candidates are accepted + only if their ``pending_request_info_events`` cover ``request_ids``. + + An incoming restored checkpoint id is excluded so a failed pause save after + resume cannot re-advertise pre-response state. + + Args: + request_ids: Pending request_info ids that must be present on the checkpoint. + checkpoint_storage: Optional storage override for this lookup. + known_checkpoint_id: Fallback id (for example a cold-resume short-circuit). + baseline_checkpoint_id: Optional override for the pre-run runner id. When omitted, + uses the baseline captured by the most recent ``run()``. Pass ``None`` + explicitly to treat any current runner id as newly advanced. + + Returns: + A checkpoint id suitable for durable resume, or ``None`` when none is safe. + """ + ids = {str(request_id) for request_id in request_ids if request_id} + if not ids: + return None + + if baseline_checkpoint_id is _MISSING: + baseline_checkpoint_id = self._run_baseline_checkpoint_id + + # Prefer explicit storage; otherwise load via public RunnerContext APIs + # (Protocol has no private `_get_effective_checkpoint_storage`). + storage = checkpoint_storage + use_context_storage = storage is None and self._runner.context.has_checkpointing() + + current = self.get_last_checkpoint_id() + excluded: set[str] = set() + if self._restored_checkpoint_id is not None: + excluded.add(self._restored_checkpoint_id) + + # Run-scoped: do not advertise a pre-run leftover when this run did not persist. + runner_candidate: str | None = None + if current is not None and current != baseline_checkpoint_id and current not in excluded: + runner_candidate = current + + candidates: list[str] = [] + if runner_candidate is not None: + candidates.append(runner_candidate) + if ( + known_checkpoint_id is not None + and known_checkpoint_id not in candidates + and known_checkpoint_id not in excluded + ): + candidates.append(str(known_checkpoint_id)) + + if storage is None and not use_context_storage: + # Without storage we cannot prove coverage; only advertise a run-scoped runner id. + return runner_candidate + + for candidate in candidates: + try: + if storage is not None: + checkpoint = await storage.load(candidate) + else: + checkpoint = await self._runner.context.load_checkpoint(candidate) + except Exception: # pragma: no cover - storage/type drift + logger.debug("Could not load pause checkpoint candidate %s", candidate, exc_info=True) + continue + if checkpoint is None: + continue + pending = checkpoint.pending_request_info_events or {} + if ids.issubset({str(key) for key in dict(pending)}): + return candidate + return None + async def cancel_pending_requests( self, request_ids: Collection[str], diff --git a/python/packages/core/agent_framework/observability.py b/python/packages/core/agent_framework/observability.py index 9ecc90d06c8..fc911223ef9 100644 --- a/python/packages/core/agent_framework/observability.py +++ b/python/packages/core/agent_framework/observability.py @@ -1556,6 +1556,7 @@ def disable_instrumentation() -> None: def enable_instrumentation( *, enable_sensitive_data: bool | None = None, + enable_message_events: bool | None = None, force: bool = False, ) -> None: """Enable instrumentation for Microsoft Agent Framework. @@ -1568,9 +1569,17 @@ def enable_instrumentation( Keyword Args: enable_sensitive_data: Enable OpenTelemetry sensitive events. Overrides the environment variable ENABLE_SENSITIVE_DATA if set. Default is None. + enable_message_events: Emit the baseline v1.36.0 GenAI message events for + model invocation when sensitive data capture is enabled. Explicit values + override the current setting, including values from ENABLE_MESSAGE_EVENTS. + Default is None, which preserves the current setting without re-reading + the environment. Does not affect experimental message span attributes. force: When True, clears any sticky disable previously set by ``disable_instrumentation()`` before enabling. Without it, calls are no-ops if instrumentation has been explicitly disabled. + + Note: + This function does not configure or replace OpenTelemetry providers or exporters. """ global OBSERVABILITY_SETTINGS if OBSERVABILITY_SETTINGS._user_disabled and not force: # type: ignore[reportPrivateUsage] @@ -1587,6 +1596,8 @@ def enable_instrumentation( else: # Re-read from current environment in case env vars were set after import (e.g. load_dotenv()) OBSERVABILITY_SETTINGS.enable_sensitive_data = _read_bool_env("ENABLE_SENSITIVE_DATA") + if enable_message_events is not None: + OBSERVABILITY_SETTINGS.enable_message_events = enable_message_events def configure_otel_providers( diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml index d2280bd94da..756b4a7522a 100644 --- a/python/packages/core/pyproject.toml +++ b/python/packages/core/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ "pydantic>=2,<3", "python-dotenv>=1,<2", "opentelemetry-api>=1.39.0,<2", + "pyyaml>=6.0,<7.0", ] [project.optional-dependencies] @@ -71,6 +72,7 @@ dev = [ "azure-ai-agentserver-core>=2.0.0b7,<3", "graphviz>=0.20,<1", "opentelemetry-exporter-otlp>=1.39.0,<2", + "types-pyyaml>=6.0,<7.0", ] [tool.uv] diff --git a/python/packages/core/tests/core/test_compaction.py b/python/packages/core/tests/core/test_compaction.py index 8e48faa4bf4..c02eb5364b4 100644 --- a/python/packages/core/tests/core/test_compaction.py +++ b/python/packages/core/tests/core/test_compaction.py @@ -435,7 +435,106 @@ def test_annotate_message_groups_with_tokenizer_adds_token_counts() -> None: assert isinstance(_token_count(messages[1]), int) +def test_annotate_message_groups_token_counts_ignore_protected_data() -> None: + encrypted_blob = "A" * 4096 + with_payload = Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_1", text="", protected_data=encrypted_blob), + Content.from_text("final answer"), + ], + ) + without_payload = Message( + role="assistant", + contents=[ + Content.from_text_reasoning(id="rs_1", text=""), + Content.from_text("final answer"), + ], + ) + + annotate_message_groups([with_payload], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([without_payload], tokenizer=CharacterEstimatorTokenizer()) + + assert _token_count(with_payload) == _token_count(without_payload) + + +def test_annotate_message_groups_token_counts_ignore_encrypted_content_in_additional_properties() -> None: + encrypted_blob = "A" * 4096 + with_blob = Content.from_text_reasoning(id="rs_1", text="") + with_blob.additional_properties["encrypted_content"] = encrypted_blob + without_blob = Content.from_text_reasoning(id="rs_1", text="") + + with_message = Message(role="assistant", contents=[with_blob, Content.from_text("final answer")]) + without_message = Message(role="assistant", contents=[without_blob, Content.from_text("final answer")]) + + annotate_message_groups([with_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([without_message], tokenizer=CharacterEstimatorTokenizer()) + + assert _token_count(with_message) == _token_count(without_message) + + +def test_annotate_message_groups_token_counts_keep_clear_text_reasoning_details() -> None: + import json + + clear_text = "visible reasoning " * 128 + with_details = Content.from_text_reasoning( + id="rs_1", + text="", + protected_data=json.dumps([{"type": "reasoning.text", "text": clear_text, "encrypted_content": "A" * 4096}]), + ) + summary_only = Content.from_text_reasoning( + id="rs_1", + text="", + protected_data=json.dumps([{"type": "reasoning.text", "text": clear_text}]), + ) + without_details = Content.from_text_reasoning(id="rs_1", text="") + final_answer = Content.from_text("final answer") + + with_message = Message(role="assistant", contents=[with_details, final_answer]) + summary_only_message = Message(role="assistant", contents=[summary_only, final_answer]) + bare_message = Message(role="assistant", contents=[without_details, final_answer]) + + annotate_message_groups([with_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([summary_only_message], tokenizer=CharacterEstimatorTokenizer()) + annotate_message_groups([bare_message], tokenizer=CharacterEstimatorTokenizer()) + + # The opaque encrypted member is excluded while the clear-text reasoning + # the provider replays stays counted. + with_count = _token_count(with_message) + summary_count = _token_count(summary_only_message) + bare_count = _token_count(bare_message) + assert with_count is not None + assert summary_count is not None + assert bare_count is not None + assert with_count == summary_count + assert with_count > bare_count + + +def test_annotate_token_counts_recomputes_counts_from_stale_serialization_basis() -> None: + + from agent_framework._compaction import TOKEN_COUNT_BASIS_VERSION, _write_group_annotation, annotate_token_counts + + message = Message(role="assistant", contents=[Content.from_text("final answer " * 64)]) + _write_group_annotation(message, group_id="g0", kind="assistant_text", index=0, has_reasoning=False) + annotation = message.additional_properties[GROUP_ANNOTATION_KEY] + assert isinstance(annotation, dict) + # Simulate a count cached by a pre-fix serialization basis: value present, + # no basis stamp (annotations survive to_dict/from_dict round-trips). + annotation[GROUP_TOKEN_COUNT_KEY] = 1 + annotation.pop("token_count_basis", None) + assert _token_count(message) == 1 + + annotate_token_counts([message], tokenizer=CharacterEstimatorTokenizer()) + + annotation = message.additional_properties[GROUP_ANNOTATION_KEY] + assert isinstance(annotation, dict) + assert annotation[GROUP_TOKEN_COUNT_KEY] != 1 + assert annotation["token_count_basis"] == TOKEN_COUNT_BASIS_VERSION + + def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() -> None: + from agent_framework._compaction import TOKEN_COUNT_BASIS_VERSION + tokenizer = CharacterEstimatorTokenizer() messages = [_assistant_function_call("c3")] annotate_message_groups(messages) @@ -445,6 +544,10 @@ def test_extend_compaction_messages_preserves_existing_annotations_and_tokens() annotation = messages[0].additional_properties.get(GROUP_ANNOTATION_KEY) if isinstance(annotation, dict): annotation[GROUP_TOKEN_COUNT_KEY] = old_token_count + # Counts cached under the current serialization basis are preserved + # across extensions; counts without the current basis stamp are stale + # and get recomputed (see test_annotate_token_counts_recomputes_...). + annotation["token_count_basis"] = TOKEN_COUNT_BASIS_VERSION extend_compaction_messages(messages, [_tool_result("c3", "ok")], tokenizer=tokenizer) diff --git a/python/packages/core/tests/core/test_function_invocation_logic.py b/python/packages/core/tests/core/test_function_invocation_logic.py index 0fcd178bb2a..9dea58bf4c2 100644 --- a/python/packages/core/tests/core/test_function_invocation_logic.py +++ b/python/packages/core/tests/core/test_function_invocation_logic.py @@ -3463,6 +3463,243 @@ def approval_func(value: str) -> str: assert approval_arguments == ["expected"] +@pytest.mark.parametrize("include_approval_response", [False, True], ids=["zero-responses", "approval-only"]) +async def test_stateless_split_mixed_batch_rejects_incomplete_replay_before_execution( + chat_client_base: SupportsChatGetResponse, + include_approval_response: bool, +) -> None: + """A split stateless mixed batch cannot execute until every response arrives.""" + from agent_framework import FunctionTool + + calls = 0 + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal calls + calls += 1 + return "approved" + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + host_func = FunctionTool(name="host_func", func=None, description="Handled by the caller") + replay_contents = ( + [approval_request.to_function_approval_response(approved=True)] + if include_approval_response + else [Content.from_text("unrelated follow-up")] + ) + messages = [ + Message(role="assistant", contents=[approval_request]), + Message(role="tool", contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message(role="user", contents=replay_contents), + ] + + with pytest.raises( + RuntimeError, + match="A mixed function-call batch requires responses for every approval and Host-owned request", + ): + await chat_client_base.get_response( + messages, + options={"tools": [approval_func, host_func]}, + ) + + assert calls == 0 + + +@pytest.mark.parametrize("metadata_role", ["tool", "user"]) +def test_stateless_mixed_batch_across_message_roles_requires_complete_responses(metadata_role: str) -> None: + """Intervening message roles do not split a mixed batch or change response order.""" + from agent_framework._tools import _stateless_mixed_pause_batch_status + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + approval_response = approval_request.to_function_approval_response(approved=True) + partial_messages = [ + Message(role="assistant", contents=[approval_request]), + Message(role=metadata_role, contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message(role="user", contents=[approval_response]), + ] + + incomplete, _ = _stateless_mixed_pause_batch_status(partial_messages) + + assert incomplete is True + + host_result = Content.from_function_result(call_id="host", result="host result") + host_result.id = "host-occurrence" + complete_messages = [ + Message(role="assistant", contents=[approval_request]), + Message(role=metadata_role, contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message(role="user", contents=[host_result, approval_response]), + ] + + incomplete, host_result_ids = _stateless_mixed_pause_batch_status(complete_messages) + + assert incomplete is False + assert [content.type for content in complete_messages[-1].contents] == [ + "function_approval_response", + "function_result", + ] + assert host_result_ids == {id(complete_messages[-1].contents[1])} + + +async def test_later_standalone_request_does_not_hide_incomplete_stateless_mixed_batch( + chat_client_base: SupportsChatGetResponse, +) -> None: + """A newer standalone request cannot let an older mixed approval execute early.""" + from agent_framework import FunctionTool + + calls = 0 + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal calls + calls += 1 + return "approved" + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + later_call = Content.from_function_call( + call_id="later", + name="approval_func", + arguments={}, + id="later-occurrence", + ) + later_request = Content.from_function_approval_request( + id="later-occurrence", + function_call=later_call, + ) + host_func = FunctionTool(name="host_func", func=None, description="Handled by the caller") + messages = [ + Message(role="user", contents=["run mixed"]), + Message(role="assistant", contents=[approval_request, host_request]), + Message(role="user", contents=[approval_request.to_function_approval_response(approved=True)]), + Message(role="assistant", contents=[later_request]), + Message(role="user", contents=["unrelated follow-up"]), + ] + + with pytest.raises( + RuntimeError, + match="A mixed function-call batch requires responses for every approval and Host-owned request", + ): + await chat_client_base.get_response( + messages, + options={"tools": [approval_func, host_func]}, + ) + + assert calls == 0 + + +async def test_completed_split_stateless_mixed_batch_is_inert_on_later_turn( + chat_client_base: SupportsChatGetResponse, +) -> None: + """A terminal result prevents completed split mixed approvals from replaying.""" + from agent_framework import FunctionTool + + calls = 0 + + @tool(name="approval_func", approval_mode="always_require") + def approval_func() -> str: + nonlocal calls + calls += 1 + return "approved" + + approval_call = Content.from_function_call( + call_id="approval", + name="approval_func", + arguments={}, + id="approval-occurrence", + ) + approval_request = Content.from_function_approval_request( + id="approval-occurrence", + function_call=approval_call, + ) + host_request = Content.from_function_call( + call_id="host", + name="host_func", + arguments={}, + id="host-occurrence", + ) + host_request.user_input_request = True + host_result = Content.from_function_result(call_id="host", result="host result") + host_result.id = "host-occurrence" + approval_result = Content.from_function_result(call_id="approval", result="approved") + host_func = FunctionTool(name="host_func", func=None, description="Handled by the caller") + chat_client_base.run_responses = [ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute] + ChatResponse(messages=Message(role="assistant", contents=["later response"])), + ] + messages = [ + Message(role="user", contents=["run mixed"]), + Message(role="assistant", contents=[approval_request]), + Message(role="tool", contents=[Content.from_text("provider metadata")]), + Message(role="assistant", contents=[host_request]), + Message( + role="user", + contents=[ + approval_request.to_function_approval_response(approved=True), + host_result, + ], + ), + Message(role="tool", contents=[approval_result]), + Message(role="assistant", contents=["done"]), + Message(role="user", contents=["later"]), + ] + + response = await chat_client_base.get_response( + messages, + options={"tools": [approval_func, host_func]}, + ) + + assert response.text == "later response" + assert calls == 0 + + def test_active_mixed_pause_ignores_historical_host_requests() -> None: """Only the session-recorded mixed batch participates in response correlation.""" from agent_framework._tools import ( diff --git a/python/packages/core/tests/core/test_harness_agent.py b/python/packages/core/tests/core/test_harness_agent.py index 6535a6fc68f..41030e3bd08 100644 --- a/python/packages/core/tests/core/test_harness_agent.py +++ b/python/packages/core/tests/core/test_harness_agent.py @@ -24,6 +24,7 @@ FileAccessProvider, FileMemoryProvider, FileSystemAgentFileStore, + FunctionTool, InMemoryAgentFileStore, InMemoryHistoryProvider, Message, @@ -33,6 +34,9 @@ SkillsProvider, TodoProvider, create_harness_agent, + get_agent_mode, + set_agent_mode, + tool, ) from agent_framework._harness._agent import DEFAULT_HARNESS_INSTRUCTIONS, _assemble_instructions from agent_framework._harness._mode import AgentModeProvider @@ -133,6 +137,56 @@ def test_create_harness_agent_disable_mode() -> None: assert AgentModeProvider not in provider_types +async def test_create_harness_agent_with_replacement_mode_tool() -> None: + """Applications can replace a built-in mode tool without replacing provider state.""" + session = AgentSession(session_id="session-1") + mode_provider = AgentModeProvider(source_id="ui_mode", expose_mode_set=False) + + @tool + def update_mode(mode: str) -> str: + """Update the application's mode.""" + return set_agent_mode( + session, + mode, + source_id=mode_provider.source_id, + available_modes=mode_provider.available_modes, + notify=False, + ) + + agent = create_harness_agent( + client=_FakeChatClient(), + max_context_window_tokens=128_000, + max_output_tokens=16_384, + mode_provider=mode_provider, + tools=[update_mode], + agent_instructions="Use update_mode for approved mode changes.", + disable_todo=True, + disable_file_memory=True, + disable_web_search=True, + ) + _, options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Start planning"])], + ) + tools = options["tools"] + assert isinstance(tools, list) + assert [mode_tool.name for mode_tool in tools if isinstance(mode_tool, FunctionTool)] == ["update_mode", "mode_get"] + assert "Use update_mode for approved mode changes." in options["instructions"] + assert "mode_set" not in options["instructions"] + replacement = next( + mode_tool for mode_tool in tools if isinstance(mode_tool, FunctionTool) and mode_tool.name == "update_mode" + ) + await replacement.invoke(arguments={"mode": "execute"}) + assert get_agent_mode(session, source_id=mode_provider.source_id) == "execute" + + updated_context, updated_options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Continue"])], + ) + assert "You are currently operating in the execute mode." in updated_options["instructions"] + assert updated_context.context_messages.get(mode_provider.source_id, []) == [] + + def test_create_harness_agent_disable_file_memory() -> None: """disable_file_memory=True should exclude the FileMemoryProvider.""" agent = create_harness_agent( diff --git a/python/packages/core/tests/core/test_harness_mode.py b/python/packages/core/tests/core/test_harness_mode.py index c4c221f33fe..57120f1ffba 100644 --- a/python/packages/core/tests/core/test_harness_mode.py +++ b/python/packages/core/tests/core/test_harness_mode.py @@ -11,11 +11,13 @@ Agent, AgentModeProvider, AgentSession, + FunctionTool, Message, SupportsChatGetResponse, get_agent_mode, set_agent_mode, ) +from agent_framework._harness._mode import DEFAULT_MODE_INSTRUCTIONS, DEFAULT_MODE_MAP def _tool_by_name(tools: list[object], name: str) -> object: @@ -210,6 +212,132 @@ async def test_agent_mode_context_provider_updates_agent_mode( assert set_agent_mode(session, "plan", source_id=provider.source_id) == "plan" +@pytest.mark.parametrize("expose_mode_set", [True, False]) +@pytest.mark.parametrize("expose_mode_get", [True, False]) +async def test_agent_mode_provider_tool_exposure( + chat_client_base: SupportsChatGetResponse, expose_mode_set: bool, expose_mode_get: bool +) -> None: + """Tool exposure must match built-in guidance without disabling the mode workflow.""" + session = AgentSession(session_id="session-1") + provider = AgentModeProvider(expose_mode_set=expose_mode_set, expose_mode_get=expose_mode_get) + agent = Agent(client=chat_client_base, context_providers=[provider]) + + _, options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Start planning"])], + ) + tools: list[object] = options.get("tools") or [] + expected_names = [ + name for name, exposed in (("mode_set", expose_mode_set), ("mode_get", expose_mode_get)) if exposed + ] + assert [tool.name for tool in tools if isinstance(tool, FunctionTool)] == expected_names + instructions = options["instructions"] + assert isinstance(instructions, str) + assert ("mode_set" in instructions) == expose_mode_set + assert ("mode_get" in instructions) == expose_mode_get + assert "### Mandatory Mode based Workflow" in instructions + assert "get user approval before proceeding" in instructions + assert "You are currently operating in the plan mode." in instructions + assert get_agent_mode(session) == "plan" + if not expose_mode_set: + assert "only after the mode has changed" in instructions + for mode_tool in tools: + assert isinstance(mode_tool, FunctionTool) + assert mode_tool.approval_mode == "never_require" + if mode_tool.name == "mode_set": + result = await mode_tool.invoke(arguments={"mode": "execute"}) + assert result[0].text is not None + assert json.loads(result[0].text) == {"mode": "execute", "message": "Mode changed to 'execute'."} + else: + result = await mode_tool.invoke() + assert result[0].text is not None + assert json.loads(result[0].text) == {"mode": "execute" if expose_mode_set else "plan"} + + +@pytest.mark.parametrize("instructions", [None, ""]) +def test_agent_mode_provider_preserves_default_instructions(instructions: str | None) -> None: + """Suppression on one provider must not alter another provider's defaults.""" + AgentModeProvider(expose_mode_set=False, expose_mode_get=False) + provider = AgentModeProvider(instructions=instructions) + mode_lines = "".join(f"#### {name}\n\n{text}\n\n" for name, text in DEFAULT_MODE_MAP.items()) + expected = DEFAULT_MODE_INSTRUCTIONS.replace("{available_modes}", mode_lines).replace("{current_mode}", "plan") + assert provider._build_instructions("plan") == expected # pyright: ignore[reportPrivateUsage] + + +@pytest.mark.parametrize("custom_instructions", [True, False]) +@pytest.mark.parametrize("custom_modes", [True, False]) +def test_agent_mode_provider_preserves_custom_instructions(custom_instructions: bool, custom_modes: bool) -> None: + """Only built-in guidance should be adapted; caller text still expands placeholders.""" + mode_text = "Custom mode_get and mode_set guidance for {current_mode}." + mode_map = {"draft": mode_text} if custom_modes else None + instructions = "Use update_mode, not mode_set or mode_get. {current_mode}\n{available_modes}" + provider = AgentModeProvider( + expose_mode_set=False, + expose_mode_get=False, + instructions=instructions if custom_instructions else None, + mode_instructions=mode_map, + ) + current_mode = "draft" if custom_modes else "plan" + rendered = provider._build_instructions(current_mode) # pyright: ignore[reportPrivateUsage] + if custom_modes: + assert mode_map == {"draft": mode_text} + assert f"Custom mode_get and mode_set guidance for {current_mode}." in rendered + if custom_instructions: + assert rendered.startswith(f"Use update_mode, not mode_set or mode_get. {current_mode}\n") + assert "{available_modes}" not in rendered + assert "{current_mode}" not in rendered + if not custom_modes and not custom_instructions: + assert "mode_set" not in rendered + assert "mode_get" not in rendered + + +async def test_agent_mode_provider_hidden_tools_preserve_state_and_notifications( + chat_client_base: SupportsChatGetResponse, +) -> None: + """State and one-shot external notifications must survive even with no mode tools.""" + session = AgentSession(session_id="session-1") + provider = AgentModeProvider( + source_id="ui_mode", default_mode="execute", expose_mode_set=False, expose_mode_get=False + ) + agent = Agent(client=chat_client_base, context_providers=[provider]) + _, first_options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Start"])], + ) + assert not first_options.get("tools") + assert "You are currently operating in the execute mode." in first_options["instructions"] + assert get_agent_mode(session, source_id=provider.source_id) == "execute" + set_agent_mode(session, "plan", source_id=provider.source_id) + + changed_context, changed_options = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Continue"])], + ) + assert "You are currently operating in the plan mode." in changed_options["instructions"] + notifications = changed_context.context_messages.get(provider.source_id, []) + assert len(notifications) == 1 + assert notifications[0].role == "user" + assert 'from "execute" to "plan"' in notifications[0].text + assert "previous_mode_for_notification" not in session.state[provider.source_id] + + set_agent_mode(session, "plan", source_id=provider.source_id) + unchanged_context, _ = await agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Continue planning"])], + ) + assert unchanged_context.context_messages.get(provider.source_id, []) == [] + + reconfigured_provider = AgentModeProvider(source_id=provider.source_id, default_mode="execute") + reconfigured_agent = Agent(client=chat_client_base, context_providers=[reconfigured_provider]) + _, restored_options = await reconfigured_agent._prepare_session_and_messages( # pyright: ignore[reportPrivateUsage] + session=session, + input_messages=[Message(role="user", contents=["Status"])], + ) + assert "You are currently operating in the plan mode." in restored_options["instructions"] + assert get_agent_mode(session, source_id=provider.source_id) == "plan" + assert DEFAULT_MODE_SOURCE_ID not in session.state + + def test_default_mode_falls_back_to_first_available_mode() -> None: """When ``default_mode`` is omitted, helpers and provider should use the first configured mode.""" session = AgentSession(session_id="session-1") @@ -251,6 +379,29 @@ def test_set_agent_mode_no_op_does_not_record_previous_mode() -> None: assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID] +def test_set_agent_mode_can_skip_external_change_notification() -> None: + """Agent-invoked replacement tools should be able to avoid a redundant notification.""" + session = AgentSession(session_id="session-1") + set_agent_mode(session, "plan") + set_agent_mode(session, "execute", notify=False) + + assert get_agent_mode(session) == "execute" + assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID] + + +def test_set_agent_mode_without_notification_clears_pending_notification() -> None: + """An agent-observed update should replace pending external transition context.""" + session = AgentSession(session_id="session-1") + set_agent_mode(session, "plan") + set_agent_mode(session, "execute") + assert session.state[DEFAULT_MODE_SOURCE_ID]["previous_mode_for_notification"] == "plan" + + set_agent_mode(session, "plan", notify=False) + + assert get_agent_mode(session) == "plan" + assert "previous_mode_for_notification" not in session.state[DEFAULT_MODE_SOURCE_ID] + + async def test_agent_mode_provider_injects_user_message_after_external_change( chat_client_base: SupportsChatGetResponse, ) -> None: diff --git a/python/packages/core/tests/core/test_harness_tool_approval.py b/python/packages/core/tests/core/test_harness_tool_approval.py index 95f65e5a467..b3b88359215 100644 --- a/python/packages/core/tests/core/test_harness_tool_approval.py +++ b/python/packages/core/tests/core/test_harness_tool_approval.py @@ -24,6 +24,7 @@ ChatResponse, ChatResponseUpdate, Content, + ContextProvider, FileHistoryProvider, FunctionInvocationContext, FunctionMiddleware, @@ -31,6 +32,8 @@ InMemoryHistoryProvider, Message, ResponseStream, + SessionContext, + SupportsAgentRun, ToolApprovalMiddleware, ToolApprovalState, create_always_approve_tool_response, @@ -134,6 +137,136 @@ def safe_tool() -> str: ] +async def test_sessionless_approval_resume_with_context_provider_uses_message_authority( + chat_client_base: MockBaseChatClient, +) -> None: + """A framework-created context-provider session must not become approval authority.""" + calls: list[str] = [] + + class NoOpContextProvider(ContextProvider): + pass + + @tool(name="approved_tool", approval_mode="always_require") + def approved_tool() -> str: + calls.append("approved") + return "approved" + + agent = Agent( + client=chat_client_base, + tools=[approved_tool], + context_providers=[NoOpContextProvider("context")], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="approved-call", + name="approved_tool", + arguments="{}", + id="approved-occurrence", + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + first = await agent.run("request approval") + approval_request = first.user_input_requests[0] + resumed = await agent.run([ + Message(role="user", contents=["request approval"]), + *first.messages, + Message(role="user", contents=[approval_request.to_function_approval_response(True)]), + ]) + + assert calls == ["approved"] + assert [(message.role, [content.type for content in message.contents]) for message in resumed.messages] == [ + ("tool", ["function_result"]), + ("assistant", ["text"]), + ] + + +async def test_framework_created_session_becomes_authoritative_when_reused_by_caller( + chat_client_base: MockBaseChatClient, +) -> None: + """A retained framework session must fail closed when later supplied explicitly.""" + calls: list[int] = [] + retained_sessions: list[AgentSession] = [] + + class RetainingContextProvider(ContextProvider): + async def before_run( + self, + *, + agent: SupportsAgentRun, + session: AgentSession, + context: SessionContext, + state: dict[str, Any], + ) -> None: + del agent, context, state + retained_sessions.append(session) + + @tool(name="approved_tool", approval_mode="always_require") + def approved_tool(value: int) -> str: + calls.append(value) + return "approved" + + agent = Agent( + client=chat_client_base, + tools=[approved_tool], + context_providers=[RetainingContextProvider("context")], + ) + chat_client_base.run_responses = [ + ChatResponse( + messages=Message( + role="assistant", + contents=[ + Content.from_function_call( + call_id="approved-call", + name="approved_tool", + arguments={"value": 1}, + id="approved-occurrence", + ) + ], + ) + ), + ChatResponse(messages=Message(role="assistant", contents=["done"])), + ] + + first = await agent.run("request approval") + assert calls == [] + assert "tool_approval" not in retained_sessions[0].state + approval_request = first.user_input_requests[0] + original_call = _function_call(approval_request) + assert original_call.call_id is not None + assert original_call.name is not None + assert approval_request.id is not None + tampered_call = Content.from_function_call( + call_id=original_call.call_id, + name=original_call.name, + arguments={"value": 999}, + id=original_call.id, + ) + tampered_response = Content.from_function_approval_response( + approved=True, + id=approval_request.id, + function_call=tampered_call, + ) + + await agent.run( + [ + Message(role="user", contents=["request approval"]), + *first.messages, + Message(role="user", contents=[tampered_response]), + ], + session=retained_sessions[0], + ) + + assert calls == [] + assert retained_sessions[1] is retained_sessions[0] + + class _StringApprovalValue(str, Enum): ALPHA = "alpha" @@ -2116,6 +2249,53 @@ def second_guarded() -> str: assert state["queued_approval_requests"] == [] +@pytest.mark.parametrize("approval_first", [True, False], ids=["approval-first", "host-first"]) +async def test_tool_approval_middleware_preserves_split_streamed_mixed_batch_order( + chat_client_base: MockBaseChatClient, + approval_first: bool, +) -> None: + """Separately streamed approval and Host pauses remain one ordered batch.""" + from agent_framework import FunctionTool + from agent_framework._middleware import AgentContext + + @tool(name="guarded", approval_mode="always_require") + def guarded() -> str: + return "guarded" + + host_tool = FunctionTool(name="host_read", func=None, description="Handled by the caller") + agent = Agent(client=chat_client_base, tools=[guarded, host_tool]) + approval_call = Content.from_function_call(call_id="guarded", name="guarded", arguments={}) + approval_request = Content.from_function_approval_request(id="guarded", function_call=approval_call) + host_request = Content.from_function_call(call_id="host", name="host_read", arguments={}) + host_request.user_input_request = True + ordered_requests = [approval_request, host_request] if approval_first else [host_request, approval_request] + context = AgentContext( + agent=agent, + messages=[], + session=AgentSession(), + stream=True, + ) + middleware = ToolApprovalMiddleware() + state = ToolApprovalState() + + async def inner_stream(): + for request in ordered_requests: + yield AgentResponseUpdate(role="assistant", contents=[request]) + + async def call_next() -> None: + context.result = ResponseStream(inner_stream(), finalizer=AgentResponse.from_updates) + + response_stream = middleware._process_stream( # pyright: ignore[reportPrivateUsage] + context, + call_next, + state, + ) + updates = [update async for update in response_stream] + + assert [content for update in updates for content in update.user_input_requests] == ordered_requests + assert state.queued_approval_requests == [] + + @pytest.mark.parametrize("streaming", [False, True], ids=["non-streaming", "streaming"]) async def test_tool_approval_middleware_policy_approval_reclassifies_host_tool( chat_client_base: MockBaseChatClient, diff --git a/python/packages/core/tests/core/test_mcp_skills.py b/python/packages/core/tests/core/test_mcp_skills.py index 20cb8db5857..85638a036d0 100644 --- a/python/packages/core/tests/core/test_mcp_skills.py +++ b/python/packages/core/tests/core/test_mcp_skills.py @@ -757,6 +757,161 @@ def _archive_client(index_json: str, archive_url: str, archive_bytes: bytes, mim class TestMCPSkillsSourceArchive: """Tests for archive-type skill discovery via MCPSkillsSource (in-memory).""" + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + "fields", + ( + "description: First\ndescription: >-\n Second", + "description: |-\n First\nDescription: Second", + "description: Valid\nmetadata:\n author: First\nmetadata:\n author: Second", + "description: Valid\nAllowed-Tools: read", + 'description: Valid\nallowed-tools: read\n"allowed-tools": other', + "description: Valid\n'allowed-tools': other\nallowed-tools: read", + 'description: Valid\n"description": Other text', + "'description': Other text\ndescription: Valid", + 'description: Valid\n"Allowed-Tools": other', + 'description: Valid\nmetadata:\n author: First\n"metadata":\n author: Second', + "description: Valid\nlicense: MIT\n'license':", + 'description: Valid\ncompatibility: Any runtime\n"compatibility": Other runtime', + 'description: Valid\n"name": packaged-skill', + 'description: First\n"descrip\\u0074ion": Second', + 'description: First\n"\\x44escription": Second', + "description: Valid\nmetadata: [unterminated", + 'description: "\\U00110000"', + 'description: "\\UFFFFFFFF"', + 'description: "\\uD800"', + 'description: "\\uDFFF"', + ), + ) + async def test_ambiguous_archive_frontmatter_is_skipped(self, fields: str, newline: str) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = f"---\nname: packaged-skill\n{fields}\n---\nBody." + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert skills == [] + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + async def test_archive_yaml_escapes_and_invalid_metadata( + self, newline: str, caplog: pytest.LogCaptureFixture + ) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = ( + '---\nname: packaged-skill\n"\\x64escription": "Read\\nfiles"\n' + '"\\u006detadata": {author: First, "\\u0061uthor": Second, invalid: [value]}\n---\nBody.' + ) + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.description == "Read\nfiles" + assert skills[0].frontmatter.metadata == {"author": "First"} + assert len(caplog.records) == 2 + assert all(record.levelname == "WARNING" for record in caplog.records) + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + async def test_archive_value_on_indented_next_line_is_preserved(self, newline: str) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = "---\nname: packaged-skill\ndescription:\n 'Read files'\nlicense: MIT\n---\nBody." + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.description == "Read files" + assert skills[0].frontmatter.license == "MIT" + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("field", ("license", "compatibility", "allowed-tools")) + async def test_archive_empty_optional_scalar_remains_none(self, field: str, newline: str) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = f"---\nname: packaged-skill\ndescription: Read files\n{field}: \n---\nBody." + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.description == "Read files" + assert skills[0].frontmatter.license is None + assert skills[0].frontmatter.compatibility is None + assert skills[0].frontmatter.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("second_key", ("author", "Author")) + @pytest.mark.parametrize("quote", ("", "'", '"')) + async def test_archive_duplicate_metadata_keeps_first_value_and_warns( + self, newline: str, second_key: str, quote: str, caplog: pytest.LogCaptureFixture + ) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = ( + f"---\n{quote}name{quote}: packaged-skill\n{quote}description{quote}: >-\n Read\n files\n" + f"{quote}compatibility{quote}:\n Any runtime\n{quote}allowed-tools{quote}: 'read'\n" + f"{quote}metadata{quote}:\n" + f" author:\n 'First'\n {second_key}: Second\n author: Third\n version: '1.0'\n" + f"{quote}license{quote}: |-\n MIT\n License\n---\nBody." + ) + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + expected = {"author": "First", "version": "1.0"} + if second_key == "Author": + expected["Author"] = "Second" + assert skills[0].frontmatter.metadata == expected + assert skills[0].frontmatter.description == "Read files" + assert skills[0].frontmatter.license == "MIT\nLicense" + assert skills[0].frontmatter.compatibility == "Any runtime" + assert skills[0].frontmatter.allowed_tools == "read" + assert "Body." in await skills[0].get_content() + assert len(caplog.records) == (2 if second_key == "author" else 1) + assert all(record.levelname == "WARNING" for record in caplog.records) + assert all( + "duplicate metadata key 'author'; keeping the first value" in record.getMessage() + for record in caplog.records + ) + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("First", "First"), + ("\n 'First'", "First"), + ("|-\n First\n Second", "First\nSecond"), + (">-\n First\n Second", "First Second"), + ), + ) + async def test_archive_yaml_metadata_scalar_formats( + self, newline: str, value: str, expected: str, caplog: pytest.LogCaptureFixture + ) -> None: + url = "skill://archives/packaged-skill.zip" + index = _make_archive_index("packaged-skill", url) + content = ( + "---\nname: packaged-skill\ndescription: Read files\n" + f"metadata:\n author: {value}\n version: '1.0'\n---\nBody." + ) + archive = _make_zip({"SKILL.md": content.replace("\n", newline).encode()}) + client = _archive_client(index, url, archive, "application/zip") + + skills = await MCPSkillsSource(client=client).get_skills(_SOURCE_CTX) + + assert len(skills) == 1 + assert skills[0].frontmatter.metadata == {"author": expected, "version": "1.0"} + assert "Body." in await skills[0].get_content() + assert not caplog.records + async def test_zip_archive_discovered_as_file_skill(self) -> None: from agent_framework import FileSkill diff --git a/python/packages/core/tests/core/test_observability.py b/python/packages/core/tests/core/test_observability.py index 0dcafce37f8..fb4a5b1d38c 100644 --- a/python/packages/core/tests/core/test_observability.py +++ b/python/packages/core/tests/core/test_observability.py @@ -1914,6 +1914,137 @@ def test_enable_instrumentation_explicit_param_overrides_env(monkeypatch): assert observability.OBSERVABILITY_SETTINGS.enable_sensitive_data is False +@pytest.mark.parametrize("message_events", [False, True]) +def test_enable_instrumentation_message_events_overrides_env(monkeypatch, message_events: bool) -> None: + import agent_framework.observability as observability + + monkeypatch.setenv("ENABLE_MESSAGE_EVENTS", str(not message_events)) + monkeypatch.setenv("ENABLE_SENSITIVE_DATA", "false") + settings = observability.ObservabilitySettings(enable_instrumentation=False) + monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", settings) + assert settings.enable_message_events is not message_events + + observability.enable_instrumentation(enable_message_events=message_events) + + assert settings.enable_instrumentation is True + assert settings.enable_message_events is message_events + assert settings.enable_sensitive_data is False + + +@pytest.mark.parametrize("message_events", [False, True]) +@pytest.mark.parametrize("source", ["environment", "providers", "instrumentation"]) +def test_enable_instrumentation_preserves_message_events(monkeypatch, message_events: bool, source: str) -> None: + import agent_framework.observability as observability + + monkeypatch.setenv("ENABLE_MESSAGE_EVENTS", str(message_events if source == "environment" else not message_events)) + settings = observability.ObservabilitySettings() + monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", settings) + if source == "providers": + with patch.object(settings, "_configure"): + observability.configure_otel_providers(enable_message_events=message_events) + elif source == "instrumentation": + observability.enable_instrumentation(enable_message_events=message_events) + + for env_value in [str(not message_events), None]: + if env_value is None: + monkeypatch.delenv("ENABLE_MESSAGE_EVENTS") + else: + monkeypatch.setenv("ENABLE_MESSAGE_EVENTS", env_value) + observability.enable_instrumentation() + assert settings.enable_message_events is message_events + observability.enable_instrumentation(enable_message_events=None) + assert settings.enable_message_events is message_events + + +@pytest.mark.parametrize("current_value", [False, True]) +@pytest.mark.parametrize("message_events", [False, True, None]) +@pytest.mark.parametrize("force", [False, True]) +def test_enable_instrumentation_message_events_respects_sticky_disable( + monkeypatch, current_value: bool, message_events: bool | None, force: bool +) -> None: + import agent_framework.observability as observability + + settings = observability.ObservabilitySettings(enable_message_events=current_value) + monkeypatch.setattr(observability, "OBSERVABILITY_SETTINGS", settings) + observability.disable_instrumentation() + + observability.enable_instrumentation(enable_message_events=message_events, enable_sensitive_data=True, force=force) + + assert settings.enable_instrumentation is force + assert settings.enable_sensitive_data is force + assert settings.is_user_disabled is not force + expected = message_events if force and message_events is not None else current_value + assert settings.enable_message_events is expected + + +@pytest.mark.parametrize("is_setup", [False, True]) +def test_enable_instrumentation_message_events_preserves_providers( + monkeypatch, span_exporter, log_record_exporter, is_setup: bool +) -> None: + from opentelemetry import metrics, trace + from opentelemetry._logs import get_logger_provider + + import agent_framework.observability as observability + + settings = observability.OBSERVABILITY_SETTINGS + monkeypatch.setattr(settings, "enable_console_exporters", True) + monkeypatch.setattr(settings, "_executed_setup", is_setup) + providers = (trace.get_tracer_provider(), get_logger_provider(), metrics.get_meter_provider()) + with patch.object(settings, "_configure") as configure: + for message_events in [False, True, None]: + observability.enable_instrumentation(enable_message_events=message_events) + assert settings.is_setup is is_setup + assert settings.enable_console_exporters is True + assert trace.get_tracer_provider() is providers[0] + assert get_logger_provider() is providers[1] + assert metrics.get_meter_provider() is providers[2] + configure.assert_not_called() + + +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("message_events", [False, True]) +@pytest.mark.parametrize("enable_sensitive_data", [False, True], indirect=True) +@pytest.mark.parametrize("semconv", ["", "gen_ai_latest_experimental"]) +async def test_enable_instrumentation_controls_message_event_emission( + mock_chat_client, + span_exporter: InMemorySpanExporter, + log_record_exporter, + enable_sensitive_data: bool, + stream: bool, + message_events: bool, + semconv: str, +) -> None: + import agent_framework.observability as observability + + observability.OBSERVABILITY_SETTINGS.otel_semconv_stability_opt_in = semconv + observability.enable_instrumentation(enable_message_events=message_events) + client = mock_chat_client() + messages = [Message("user", ["Test message"])] + if stream: + response_stream = client.get_response(messages=messages, options={"model": "Test"}, stream=True) + async for _ in response_stream: + pass + await response_stream.get_final_response() + else: + await client.get_response(messages=messages, options={"model": "Test"}) + + records = [record.log_record for record in log_record_exporter.get_finished_logs()] + expected_events = [] + if message_events and enable_sensitive_data: + expected_events = [OtelAttr.USER_MESSAGE.value] + if stream: + expected_events.append(OtelAttr.CHOICE.value) + assert [record.event_name for record in records] == expected_events + if records: + assert records[0].body == {"content": "Test message"} + + spans = span_exporter.get_finished_spans() + assert len(spans) == 1 + attributes = spans[0].attributes or {} + assert (OtelAttr.INPUT_MESSAGES in attributes) is (enable_sensitive_data and bool(semconv)) + assert (OtelAttr.OUTPUT_MESSAGES in attributes) is (enable_sensitive_data and bool(semconv)) + + def test_enable_instrumentation_does_not_touch_console_exporters(monkeypatch): """Test enable_instrumentation does not modify enable_console_exporters (it is an exporter concern).""" import importlib diff --git a/python/packages/core/tests/core/test_skills.py b/python/packages/core/tests/core/test_skills.py index bf6fc571b77..f15d6d93038 100644 --- a/python/packages/core/tests/core/test_skills.py +++ b/python/packages/core/tests/core/test_skills.py @@ -525,6 +525,583 @@ def test_extra_fields_parsed(self) -> None: assert result is not None assert result.name == "test-skill" + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_yaml_escaped_names_and_scalar_aliases(self, newline: str, caplog: pytest.LogCaptureFixture) -> None: + content = ( + '---\n"\\u006eame": test-skill\n"\\x64escription": &text "Read\\nfiles"\n' + 'license: *text\n"\\U0000006detadata":\n "author": First\n' + ' "\\u0061uthor": Second\n Author: Separate\n version: 1.0\n---\nBody.' + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.name == "test-skill" + assert result.description == "Read\nfiles" + assert result.license == "Read\nfiles" + assert result.metadata == {"author": "First", "Author": "Separate", "version": "1.0"} + assert len(caplog.records) == 1 + assert "duplicate metadata key 'author'" in caplog.text + + @pytest.mark.parametrize("field", ("name", "description", "license", "compatibility", "metadata", "allowed-tools")) + @pytest.mark.parametrize(("escape", "width"), (("x", 2), ("u", 4), ("U", 8))) + @pytest.mark.parametrize("uppercase", (False, True)) + def test_escaped_root_names_obey_validation( + self, field: str, escape: str, width: int, uppercase: bool, caplog: pytest.LogCaptureFixture + ) -> None: + spelling = field.upper() if uppercase else field + encoded = f'"\\{escape}{ord(spelling[0]):0{width}x}{spelling[1:]}"' + fields = { + "name": "test-skill", + "description": "Read files", + "license": "MIT", + "compatibility": "Any runtime", + "metadata": "{}", + "allowed-tools": "read", + } + content = "---\n" + "".join(f"{key}: {value}\n" for key, value in fields.items()) + content += f"{encoded}: {fields[field]}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + expected = "incorrectly cased frontmatter field" if uppercase else "duplicate frontmatter field" + assert expected in caplog.text + + def test_flow_mapping_preserves_duplicate_metadata(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\n{name: test-skill, description: Read files, " + "metadata: {author: First, author: Second, Author: Separate, enabled: true}}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First", "Author": "Separate", "enabled": "true"} + assert len(caplog.records) == 1 + assert "duplicate metadata key 'author'" in caplog.text + + def test_flow_mapping_duplicate_root_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = '---\n{name: test-skill, description: First, "\\x64escription": Second}\n---\nBody.' + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "duplicate frontmatter field 'description'" in caplog.text + + def test_aliased_root_key_duplicate_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = "---\nname: test-skill\n? &key description\n: First\n? *key\n: Second\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "duplicate frontmatter field 'description'" in caplog.text + + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("true", "true"), + ("0123", "0123"), + ("1.20", "1.20"), + ("2026-01-02", "2026-01-02"), + ("!!str null", "null"), + ('"\\u00e9"', "\u00e9"), + ('"\\U0001F600"', "\U0001f600"), + ), + ) + def test_yaml_scalar_types_remain_text(self, value: str, expected: str) -> None: + content = ( + f"---\nname: test-skill\ndescription: {value}\nlicense: {value}\ncompatibility: {value}\n" + f"allowed-tools: {value}\nmetadata:\n value: {value}\n true: boolean key\n 1: numeric key\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == expected + assert result.license == expected + assert result.compatibility == expected + assert result.allowed_tools == expected + assert result.metadata == {"value": expected, "true": "boolean key", "1": "numeric key"} + + def test_unknown_root_values_are_not_constructed(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "unknown: !!python/object/apply:builtins.str [ignored]\n" + "unknown: &recursive {self: *recursive}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert not caplog.records + + def test_root_merge_key_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = "---\nname: test-skill\ndescription: Read files\n<<: {description: Override}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid frontmatter property name" in caplog.text + + def test_metadata_merge_key_is_skipped(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "metadata: {<<: {author: Ignored}, author: First}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First"} + assert len(caplog.records) == 1 + assert "invalid metadata key; skipping entry" in caplog.text + + @pytest.mark.parametrize("value", ("", "null", "~")) + def test_yaml_null_optional_fields_remain_unset(self, value: str, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + f"metadata: {value}\nlicense: {value}\ncompatibility: {value}\nallowed-tools: {value}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata is None + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + assert not caplog.records + + @pytest.mark.parametrize("value", ("[First, Second]", "text", "42", "!!null {key: value}")) + def test_invalid_metadata_mapping_only_warns(self, value: str, caplog: pytest.LogCaptureFixture) -> None: + content = f"---\nname: test-skill\ndescription: Read files\nmetadata: {value}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata is None + assert len(caplog.records) == 1 + assert caplog.records[0].levelname == "WARNING" + assert "invalid metadata; expected a mapping" in caplog.text + + def test_invalid_metadata_entries_are_skipped(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + " valid: Keep\n mapping: {key: value}\n sequence: [First, Second]\n missing:\n" + " tagged: !!python/name:builtins.str text\n ? [complex, key]\n : value\n" + " author: {invalid: value}\n author: First valid\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"valid": "Keep", "author": "First valid"} + assert len(caplog.records) == 6 + assert all(record.levelname == "WARNING" for record in caplog.records) + assert all("skipping entry" in record.getMessage() for record in caplog.records) + + def test_recursive_metadata_alias_is_skipped(self, caplog: pytest.LogCaptureFixture) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "metadata: &metadata {recursive: *metadata, author: First}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First"} + assert len(caplog.records) == 1 + assert "invalid metadata value for 'recursive'" in caplog.text + + @pytest.mark.parametrize("escape", ("\\uD800", "\\uDFFF", "\\uD83D\\uDE00")) + @pytest.mark.parametrize("invalid_key", (False, True)) + def test_metadata_surrogates_are_skipped( + self, escape: str, invalid_key: bool, caplog: pytest.LogCaptureFixture + ) -> None: + entry = f'"{escape}": Ignored' if invalid_key else f'author: "{escape}"' + content = f"---\nname: test-skill\ndescription: Read files\nmetadata:\n {entry}\n author: First valid\n---" + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.metadata == {"author": "First valid"} + assert len(caplog.records) == 1 + assert caplog.records[0].levelname == "WARNING" + assert ("invalid metadata key" if invalid_key else "invalid metadata value") in caplog.text + + @pytest.mark.parametrize("field", ("name", "description", "license", "compatibility", "allowed-tools")) + @pytest.mark.parametrize( + "value", + ( + "[First, Second]", + "{key: value}", + "!!null [value]", + "!!python/name:builtins.str text", + '"\\uD800"', + '"\\uDFFF"', + '"\\uD83D\\uDE00"', + ), + ) + def test_non_text_root_values_rejected(self, field: str, value: str, caplog: pytest.LogCaptureFixture) -> None: + fields = {"name": "test-skill", "description": "Read files", field: value} + content = "---\n" + "".join(f"{key}: {item}\n" for key, item in fields.items()) + "---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert f"invalid '{field}' value; expected a text scalar" in caplog.text + + @pytest.mark.parametrize("header", ("", "null", "[]", "plain text", "!!python/object/apply:builtins.str [text]")) + def test_non_mapping_frontmatter_rejected(self, header: str, caplog: pytest.LogCaptureFixture) -> None: + result = FileSkillsSource._extract_frontmatter(f"---\n{header}\n---\nBody.", "test.md") + + assert result is None + assert "must contain a YAML frontmatter mapping" in caplog.text + + @pytest.mark.parametrize( + "fields", + ( + "metadata: [unterminated", + 'license: "\\q"', + 'license: "\\U00110000"', + 'license: "\\UFFFFFFFF"', + "license: *missing", + "description: \tvalue", + ), + ) + def test_invalid_yaml_rejected(self, fields: str, caplog: pytest.LogCaptureFixture) -> None: + content = f"---\nname: test-skill\ndescription: Read files\n{fields}\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid YAML frontmatter" in caplog.text + + @pytest.mark.parametrize("key", ("[complex, key]", '"\\uD800"', '"\\uDFFF"')) + def test_invalid_root_property_name_rejected(self, key: str, caplog: pytest.LogCaptureFixture) -> None: + content = f"---\nname: test-skill\ndescription: Read files\n? {key}\n: value\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid frontmatter property name" in caplog.text + + def test_excessive_yaml_nesting_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + content = "---\nname: test-skill\ndescription: Read files\nunknown: " + "[" * 2000 + "]" * 2000 + "\n---" + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + assert "invalid YAML frontmatter" in caplog.text + + @pytest.mark.parametrize( + "field", + ("name", "description", "license", "compatibility", "metadata", "allowed-tools"), + ) + def test_duplicate_recognized_field_rejected(self, field: str) -> None: + if field == "name": + fields = "name: test-skill\nname: test-skill\ndescription: A test skill." + elif field == "description": + fields = "name: test-skill\ndescription: A test skill.\ndescription: A second description." + elif field == "metadata": + fields = ( + "name: test-skill\ndescription: A test skill.\nmetadata:\n author: first\nmetadata:\n version: 1.0" + ) + else: + fields = f"name: test-skill\ndescription: A test skill.\n{field}: first\n{field}: second" + + result = FileSkillsSource._extract_frontmatter(f"---\n{fields}\n---\nBody.", "test.md") + + assert result is None + + @pytest.mark.parametrize( + ("field", "incorrect_casing"), + ( + ("name", "Name"), + ("description", "Description"), + ("license", "License"), + ("compatibility", "Compatibility"), + ("metadata", "Metadata"), + ("allowed-tools", "Allowed-Tools"), + ), + ) + def test_incorrectly_cased_recognized_field_rejected(self, field: str, incorrect_casing: str) -> None: + if field == "name": + fields = f"{incorrect_casing}: test-skill\ndescription: A test skill." + elif field == "description": + fields = f"name: test-skill\n{incorrect_casing}: A test skill." + elif field == "metadata": + fields = f"name: test-skill\ndescription: A test skill.\n{incorrect_casing}:\n author: test" + else: + fields = f"name: test-skill\ndescription: A test skill.\n{incorrect_casing}: value" + + result = FileSkillsSource._extract_frontmatter(f"---\n{fields}\n---\nBody.", "test.md") + + assert result is None + + @pytest.mark.parametrize("field", ("name", "description", "license", "compatibility", "metadata", "allowed-tools")) + @pytest.mark.parametrize("quote", ("'", '"')) + @pytest.mark.parametrize("layout", ("uppercase", "before", "after", "both", "mixed", "empty-first", "empty-last")) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_conflicting_quoted_fields_rejected( + self, + field: str, + quote: str, + layout: str, + newline: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + values = { + "name": "test-skill", + "description": "Read files", + "license": "MIT", + "compatibility": "Any runtime", + "metadata": "\n author: First", + "allowed-tools": "read", + } + fields = "\n".join(f"{key}: {value}" for key, value in values.items() if key != field) + quoted_key = f"{quote}{field}{quote}" + quoted_field = f"{quoted_key}: {values[field]}" + bare_field = f"{field}: {values[field]}" + other_quote = '"' if quote == "'" else "'" + declarations = { + "uppercase": f"{quote}{field.upper()}{quote}: {values[field]}", + "before": f"{quoted_field}\n{bare_field}", + "after": f"{bare_field}\n{quoted_field}", + "both": f"{quoted_field}\n{quoted_field}", + "mixed": f"{quoted_field}\n{other_quote}{field}{other_quote}: {values[field]}", + "empty-first": f"{quoted_key}:\n{bare_field}", + "empty-last": f"{bare_field}\n{quoted_key}:", + } + content = f"---\n{fields}\n{declarations[layout]}\n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + diagnostic = ( + f"incorrectly cased frontmatter field '{field.upper()}'; expected '{field}'" + if layout == "uppercase" + else f"duplicate frontmatter field '{field}'" + ) + assert diagnostic in caplog.text + + @pytest.mark.parametrize("quote", ("'", '"')) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_quoted_empty_optional_fields_remain_unset(self, quote: str, newline: str) -> None: + content = ( + f"---\n{quote}metadata{quote}:\n{quote}name{quote}: test-skill\n" + f"{quote}description{quote}: Read files\n{quote}license{quote}:\n" + f"{quote}compatibility{quote}:\n{quote}allowed-tools{quote}:\n---\nBody." + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.name == "test-skill" + assert result.description == "Read files" + assert result.metadata is None + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("field", ("metadata", "license", "compatibility", "allowed-tools", "vendor-option")) + def test_empty_inline_value_does_not_consume_next_field(self, field: str, newline: str) -> None: + content = newline.join(( + "---", + f"{field}: ", + "name: test-skill", + "description: A test skill.", + "---", + "Body.", + )) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.name == "test-skill" + assert result.description == "A test skill." + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize("field", ("license", "compatibility", "allowed-tools")) + def test_empty_optional_scalar_at_end_remains_none(self, field: str, newline: str) -> None: + content = f"---\nname: test-skill\ndescription: Read files\n{field}: \n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.license is None + assert result.compatibility is None + assert result.allowed_tools is None + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + "fields", + ( + "name:", + "description:", + "metadata:\nmetadata:", + "license:\nlicense: MIT", + "license: MIT\nlicense:", + "compatibility:\ncompatibility: Any runtime", + "allowed-tools:\nallowed-tools: read", + "Metadata:", + "allowed-tools: read\nALLOWED-TOOLS:", + ), + ) + def test_empty_inline_value_does_not_bypass_key_validation(self, fields: str, newline: str) -> None: + content = f"---\n{fields}\nname: test-skill\ndescription: A test skill.\n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + + @pytest.mark.parametrize("first_value", ("First", "'First'", '"First"', "|-\n First", ">-\n First")) + @pytest.mark.parametrize("second_value", ("Second", "'Second'", '"Second"', "|-\n Second", ">-\n Second")) + @pytest.mark.parametrize("second_key", ("description", "Description", "DESCRIPTION")) + def test_duplicate_and_cased_fields_across_scalar_formats( + self, first_value: str, second_value: str, second_key: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + f"---\nname: test-skill\ndescription: {first_value}\n{second_key}: {second_value}\nlicense: MIT\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + diagnostic = ( + "duplicate frontmatter field 'description'" + if second_key == "description" + else f"incorrectly cased frontmatter field '{second_key}'; expected 'description'" + ) + assert diagnostic in caplog.text + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("key", "first_value", "second_value", "expected"), + ( + ("author", "First", "Second", "First"), + ("Author", "'First'", '"Second"', "First"), + ("author", '"First"', "'Second'", "First"), + ("author", "Same", "Same", "Same"), + ("author", "''", "Second", ""), + ("author", '""', "Second", ""), + ("author", '" First "', "Second", " First "), + ("author", "0", "Second", "0"), + ("description", "First", "Second", "First"), + ("metadata", "First", "Second", "First"), + ("vendor-key", "First", "Second", "First"), + ("vendor_key", "First", "Second", "First"), + ("author", "\n 'First'", "Second", "First"), + ("author", "First", '\n "Second"', "First"), + ("author", '"author: First"', "Second", "author: First"), + ("author", "First\n # author: Not a field", "Second", "First"), + ("author", "|-\n First\n paragraph", "Second", "First\nparagraph"), + ("author", "First", ">-\n Second\n paragraph", "First"), + ("author", ">\n First\n paragraph", "|\n Second\n paragraph", "First paragraph\n"), + ), + ) + def test_duplicate_metadata_keeps_first_value_and_warns( + self, + newline: str, + key: str, + first_value: str, + second_value: str, + expected: str, + caplog: pytest.LogCaptureFixture, + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + f" {key}: {first_value}\n other: Preserved\n {key}: {second_value}\n" + f" {key}: Third\n tail: Retained\nlicense: MIT\n---\nBody." + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.license == "MIT" + assert result.metadata == {key: expected, "other": "Preserved", "tail": "Retained"} + assert [(record.levelname, record.getMessage()) for record in caplog.records] == [ + ("WARNING", f"SKILL.md at 'test.md' contains duplicate metadata key '{key}'; keeping the first value") + ] * 2 + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("first_key", "second_key"), + ( + ("author", "Author"), + ("Author", "author"), + ("author", "AUTHOR"), + ("description", "Description"), + ("metadata", "Metadata"), + ("vendor-key", "Vendor-Key"), + ), + ) + def test_metadata_keys_are_case_sensitive( + self, newline: str, first_key: str, second_key: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + f" {first_key}: First\n {second_key}: Second\n---\nBody." + ).replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.metadata == {first_key: "First", second_key: "Second"} + assert not caplog.records + + def test_nested_metadata_keeps_first_value_and_unknown_fields_are_ignored(self) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + "# description: Not a field\n" + "metadata:\n author: First\n author: Second\n Description: Nested text\n" + " \"description\": Nested quoted text\n 'metadata': Nested metadata value\n" + "vendor-option: First\nvendor-option: Second\n" + '\'vendor-option\': Third\n"VENDOR-OPTION": Fourth\n---\n"description": Body text' + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == "Read files" + assert result.metadata == { + "author": "First", + "Description": "Nested text", + "description": "Nested quoted text", + "metadata": "Nested metadata value", + } + + @pytest.mark.parametrize("value", ("''", '""')) + def test_empty_quoted_optional_values_are_empty_strings(self, value: str) -> None: + content = ( + f"---\nname: test-skill\ndescription: Read files\n" + f"license: {value}\ncompatibility: {value}\nallowed-tools: {value}\n---\nBody." + ) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.license == "" + assert result.compatibility == "" + assert result.allowed_tools == "" + + @pytest.mark.parametrize("value", ("|", "|-", "|+", ">", ">-", ">+")) + def test_empty_block_before_another_field_rejected(self, value: str) -> None: + content = f"---\nname: test-skill\ndescription: {value}\n\nlicense: MIT\n---\nBody." + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is None + # --------------------------------------------------------------------------- # Tests: skill discovery and loading @@ -534,12 +1111,212 @@ def test_extra_fields_parsed(self) -> None: class TestDiscoverAndLoadSkills: """Tests for file skill discovery via FileSkillsSource.get_skills().""" + @pytest.mark.parametrize( + ("fields", "loads"), + ( + ( + ( + '"\\x64escription": "Read\\nfiles"\n' + 'metadata: {author: First, "\\u0061uthor": Ignored, invalid: [item]}' + ), + True, + ), + ('description: Read files\n"\\x64escription": Second', False), + ('"\\x44escription": Read files', False), + ("description: Read files\nmetadata: [unterminated", False), + ('description: "\\U00110000"', False), + ('description: "\\UFFFFFFFF"', False), + ('description: "\\uD800"', False), + ('description: "\\uDFFF"', False), + ), + ) + async def test_file_yaml_decoding_and_validation( + self, tmp_path: Path, fields: str, loads: bool, caplog: pytest.LogCaptureFixture + ) -> None: + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_text(f"---\nname: test-skill\n{fields}\n---\nBody.", encoding="utf-8") + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + if loads: + assert len(skills) == 1 + assert skills["test-skill"].frontmatter.description == "Read\nfiles" + assert skills["test-skill"].frontmatter.metadata == {"author": "First"} + assert len(caplog.records) == 2 + assert all(record.levelname == "WARNING" for record in caplog.records) + else: + assert skills == {} + assert any(record.levelname == "ERROR" for record in caplog.records) + + @pytest.mark.parametrize("metadata_first", (False, True)) + @pytest.mark.parametrize("duplicate_metadata", (False, True)) + @pytest.mark.parametrize("quote", ("", "'", '"')) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("Read files", "Read files"), + ("'Read files'", "Read files"), + ('"Read files"', "Read files"), + ('" Read files "', " Read files "), + ('"Use #tags: safely"', "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", 'Use "quotes" safely'), + ('"| not a block"', "| not a block"), + ("\n Read files", "Read files"), + ("\n\n 'Read files'", "Read files"), + ("\n >-\n Read\n files", "Read files"), + ("|\n Read\n files", "Read\nfiles\n"), + ("|-\n Read\n files", "Read\nfiles"), + ("|+\n Read\n files", "Read\nfiles\n"), + (">\n Read\n files", "Read files\n"), + (">-\n Read\n files", "Read files"), + (">+\n Read\n files", "Read files\n"), + ("|-\n\n Read\n files", "\nRead\nfiles"), + ("|-\n Read\n\n files", "Read\n\nfiles"), + ("|-\n Read\n indented\n files", "Read\n indented\nfiles"), + ("|-\n description: text\n Description: text", "description: text\nDescription: text"), + ("|-\n \"description\": text\n 'allowed-tools': text", "\"description\": text\n'allowed-tools': text"), + ("Read files # comment", "Read files"), + ('"Read\\nfiles"', "Read\nfiles"), + ("Read\n files", "Read files"), + ("|2-\n Read\n files", "Read\nfiles"), + (">2-\n Read\n files", "Read files"), + ), + ) + async def test_yaml_scalar_formats( + self, + tmp_path: Path, + newline: str, + value: str, + expected: str, + metadata_first: bool, + duplicate_metadata: bool, + quote: str, + ) -> None: + fields = ( + f"{quote}name{quote}: test-skill\n{quote}description{quote}: {value}\n" + f"{quote}license{quote}: 'MIT'\n{quote}compatibility{quote}: Any runtime\n" + f"{quote}allowed-tools{quote}: read\n" + ) + metadata = f"{quote}metadata{quote}:\n author: test\n" + (" author: Ignored\n" if duplicate_metadata else "") + content = "---\n" + (metadata + fields if metadata_first else fields + metadata) + "---\nBody." + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert len(skills) == 1 + frontmatter = skills["test-skill"].frontmatter + assert frontmatter.description == expected + assert frontmatter.license == "MIT" + assert frontmatter.compatibility == "Any runtime" + assert frontmatter.allowed_tools == "read" + assert frontmatter.metadata == {"author": "test"} + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + ("value", "expected"), + ( + ("First", "First"), + ("'First'", "First"), + ('"First"', "First"), + ('" First "', " First "), + ('"Use #tags: safely"', "Use #tags: safely"), + ("\"Use 'quotes' safely\"", "Use 'quotes' safely"), + ("'Use \"quotes\" safely'", 'Use "quotes" safely'), + ("''", ""), + ('""', ""), + ('"\\n"', "\n"), + ("'It''s fine'", "It's fine"), + ("\n First", "First"), + ("\n 'First'", "First"), + ('\n "First"', "First"), + ("First\n Second", "First Second"), + ('"First\n Second"', "First Second"), + ("'First\n Second'", "First Second"), + ("|\n First\n Second", "First\nSecond\n"), + ("|-\n First\n Second", "First\nSecond"), + ("|+\n First\n Second", "First\nSecond\n"), + (">\n First\n Second", "First Second\n"), + (">-\n First\n Second", "First Second"), + (">+\n First\n Second", "First Second\n"), + ("\n |-\n First\n Second", "First\nSecond"), + ), + ) + async def test_yaml_metadata_scalar_formats( + self, tmp_path: Path, newline: str, value: str, expected: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\n" + f"metadata:\n author: {value}\n version: '1.0'\nlicense: MIT\n---\nBody." + ) + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert len(skills) == 1 + skill = skills["test-skill"] + assert skill.frontmatter.description == "Read files" + assert skill.frontmatter.license == "MIT" + assert skill.frontmatter.metadata == {"author": expected, "version": "1.0"} + assert "Body." in await skill.get_content() + assert not caplog.records + async def test_discovers_valid_skill(self, tmp_path: Path) -> None: _write_skill(tmp_path, "my-skill") skills = await _discover_file_skills_for_test([str(tmp_path)]) assert "my-skill" in skills assert skills["my-skill"].frontmatter.name == "my-skill" + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + @pytest.mark.parametrize( + "fields", + ( + 'allowed-tools: read\n"allowed-tools": other', + "'allowed-tools': other\nallowed-tools: read", + '"Description": Other text', + "metadata:\n author: First\n'metadata':\n author: Second", + ), + ) + async def test_conflicting_quoted_root_field_prevents_discovery( + self, tmp_path: Path, fields: str, newline: str + ) -> None: + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + content = f"---\nname: test-skill\ndescription: Read files\n{fields}\n---\nBody." + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert not skills + + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + async def test_duplicate_metadata_does_not_prevent_discovery( + self, tmp_path: Path, newline: str, caplog: pytest.LogCaptureFixture + ) -> None: + content = ( + "---\nname: test-skill\ndescription: Read files\nmetadata:\n" + " author: First\n Author: Separate\n author: Ignored\n version: '1.0'\n---\nBody." + ) + skill_dir = tmp_path / "test-skill" + skill_dir.mkdir() + (skill_dir / "SKILL.md").write_bytes(content.replace("\n", newline).encode()) + + skills = await _discover_file_skills_for_test([str(tmp_path)]) + + assert len(skills) == 1 + skill = skills["test-skill"] + assert skill.frontmatter.metadata == {"author": "First", "Author": "Separate", "version": "1.0"} + assert "Body." in await skill.get_content() + assert len(caplog.records) == 1 + assert caplog.records[0].levelname == "WARNING" + assert "duplicate metadata key 'author'; keeping the first value" in caplog.text + async def test_discovers_nested_skills(self, tmp_path: Path) -> None: skills_dir = tmp_path / "skills" _write_skill(skills_dir, "skill-a") @@ -2515,11 +3292,21 @@ def test_literal_block_scalar(self) -> None: assert result is not None assert result.description == "Line one\nLine two\n" + @pytest.mark.parametrize(("indicator", "expected"), (("|-", "Read"), ("|", "Read\n"), ("|+", "Read\n\n\n"))) + @pytest.mark.parametrize("newline", ("\n", "\r\n")) + def test_terminal_block_chomping_preserves_yaml_content(self, indicator: str, expected: str, newline: str) -> None: + content = f"---\nname: test-skill\ndescription: {indicator}\n Read\n\n\n---\nBody.".replace("\n", newline) + + result = FileSkillsSource._extract_frontmatter(content, "test.md") + + assert result is not None + assert result.description == expected + def test_folded_block_scalar(self) -> None: content = "---\nname: test-skill\ndescription: >\n This is a multi-line\n description block\n---\nBody." result = FileSkillsSource._extract_frontmatter(content, "test.md") assert result is not None - assert result.description == "This is a multi-line description block" + assert result.description == "This is a multi-line description block\n" def test_literal_strip_chomping(self) -> None: content = "---\nname: test-skill\ndescription: |-\n No trailing newline\n---\nBody." @@ -2583,14 +3370,14 @@ def test_folded_block_scalar_real_skill_format(self) -> None: assert result.description == ( "Coding standards, conventions, and patterns for developing Python code in the " "Agent Framework repository. Use this when writing or modifying Python source " - "files in the python/ directory." + "files in the python/ directory.\n" ) def test_block_scalar_with_other_fields_after(self) -> None: content = "---\nname: test-skill\ndescription: >\n A folded\n description\nlicense: MIT\n---\nBody." result = FileSkillsSource._extract_frontmatter(content, "test.md") assert result is not None - assert result.description == "A folded description" + assert result.description == "A folded description\n" assert result.license == "MIT" def test_plain_value_unchanged(self) -> None: @@ -2623,14 +3410,14 @@ def test_block_scalar_on_license_field(self) -> None: ) result = FileSkillsSource._extract_frontmatter(content, "test.md") assert result is not None - assert result.license == "Custom license spanning multiple lines" + assert result.license == "Custom license spanning multiple lines\n" - def test_block_scalar_tab_indentation(self) -> None: - """Tab characters should count as indentation for block scalar continuation lines.""" + def test_block_scalar_tab_indentation_rejected(self, caplog: pytest.LogCaptureFixture) -> None: + """YAML indentation must use spaces.""" content = "---\nname: test-skill\ndescription: |\n\tTab-indented line one\n\tTab-indented line two\n---\nBody." result = FileSkillsSource._extract_frontmatter(content, "test.md") - assert result is not None - assert result.description == "Tab-indented line one\nTab-indented line two\n" + assert result is None + assert "invalid YAML frontmatter" in caplog.text def test_block_scalar_blank_line_within_block(self) -> None: """Blank lines within a block scalar should be preserved as paragraph separators.""" diff --git a/python/packages/core/tests/core/test_tools.py b/python/packages/core/tests/core/test_tools.py index 89aa73486cc..2f39e265ee0 100644 --- a/python/packages/core/tests/core/test_tools.py +++ b/python/packages/core/tests/core/test_tools.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import copy import logging import threading from typing import Annotated, Any, Literal, get_args, get_origin @@ -19,6 +20,7 @@ from agent_framework._middleware import FunctionInvocationContext from agent_framework._tools import ( _auto_invoke_function, + _format_tool_parameters, _parse_annotation, _parse_inputs, normalize_function_invocation_configuration, @@ -28,6 +30,158 @@ # region FunctionTool and tool decorator tests +def test_format_tool_parameters_compact_preserves_scalar_metadata(): + schema = { + "type": "object", + "title": "InventoryInput", + "properties": { + "partNumber": {"type": "string", "description": "Part identifier", "title": "Part Number"}, + "units": {"type": "integer", "default": 1}, + "currency": {"type": "string", "enum": ["EUR", "USD"], "default": "EUR"}, + "price": {"type": "number"}, + "available": {"type": "boolean", "default": False}, + "empty": {"type": "null", "default": None}, + }, + "required": ["partNumber", "empty"], + } + original = copy.deepcopy(schema) + + effective_format, parameters = _format_tool_parameters(schema, parameter_format="compact") + + assert effective_format == "compact" + assert parameters == { + "partNumber": {"type": "string", "required": True, "description": "Part identifier"}, + "units": {"type": "integer", "required": False, "default": 1}, + "currency": {"type": "string", "required": False, "enum": ["EUR", "USD"], "default": "EUR"}, + "price": {"type": "number", "required": False}, + "available": {"type": "boolean", "required": False, "default": False}, + "empty": {"type": "null", "required": True, "default": None}, + } + parameters["currency"]["enum"].append("GBP") + assert schema == original + + +@pytest.mark.parametrize("required", [[], ["value"]]) +def test_format_tool_parameters_does_not_infer_requiredness_from_defaults(required): + schema = { + "type": "object", + "properties": {"value": {"type": "string", "default": "default"}}, + "required": required, + } + + effective_format, parameters = _format_tool_parameters(schema, parameter_format="compact") + + assert effective_format == "compact" + assert parameters["value"]["required"] == ("value" in required) + assert parameters["value"]["default"] == "default" + + +@pytest.mark.parametrize( + "schema", + [ + {"type": "object", "properties": {}}, + {"type": "object", "properties": {}, "required": [], "title": "EmptyInput"}, + ], +) +def test_format_tool_parameters_empty_object(schema): + assert _format_tool_parameters(schema, parameter_format="compact") == ("compact", {}) + + +@pytest.mark.parametrize( + "property_schema", + [ + {"type": "object", "properties": {"nested": {"type": "string"}}, "required": ["nested"]}, + {"type": "array", "items": {"type": "integer"}}, + {"$ref": "#/$defs/Address"}, + {"anyOf": [{"type": "string"}, {"type": "null"}]}, + {"type": ["string", "null"]}, + {"type": "string", "minLength": 1}, + {"type": "number", "minimum": 0}, + {"type": "integer", "exclusiveMaximum": 10}, + {"type": "string", "pattern": "^[A-Z]+$"}, + {"type": "string", "format": "date-time"}, + {"type": "string", "const": "fixed"}, + {"type": "string", "x-custom-keyword": {"constraint": "custom"}}, + {"type": "string", "allOf": [{"maxLength": 10}]}, + {"description": "An unconstrained parameter"}, + True, + False, + ], +) +def test_format_tool_parameters_compact_falls_back_for_rich_properties(property_schema): + schema = {"type": "object", "properties": {"value": property_schema}, "required": ["value"]} + + effective_format, parameters = _format_tool_parameters(schema, parameter_format="compact") + + assert effective_format == "json" + assert parameters == schema + + +@pytest.mark.parametrize( + "extra", + [ + {"additionalProperties": False}, + {"additionalProperties": {"type": "string"}}, + {"$defs": {"Address": {"type": "string"}}}, + {"oneOf": [{"required": ["value"]}, {"required": ["other"]}]}, + {"dependentRequired": {"value": ["other"]}}, + {"patternProperties": {"^x-": {"type": "integer"}}}, + {"minProperties": 1}, + {"x-custom-keyword": "preserve"}, + {"required": ["unlisted"]}, + {"required": "value"}, + {"required": [1]}, + {"properties": []}, + {"properties": {1: {"type": "string"}}}, + {"properties": {"value": {"type": "string", "required": True}}}, + {"type": "array"}, + ], +) +def test_format_tool_parameters_compact_falls_back_for_root_constraints(extra): + schema = {"type": "object", "properties": {"value": {"type": "string"}}, **extra} + + effective_format, parameters = _format_tool_parameters(schema, parameter_format="compact") + + assert effective_format == "json" + assert parameters == schema + + +@pytest.mark.parametrize("parameter_format", ["compact", "json"]) +def test_format_tool_parameters_full_schema_result_is_detached(parameter_format): + schema = { + "type": "object", + "properties": {"value": {"$ref": "#/$defs/Value"}}, + "$defs": {"Value": {"type": "string", "enum": ["one", "two"]}}, + "required": ["value"], + "additionalProperties": False, + } + original = copy.deepcopy(schema) + + effective_format, parameters = _format_tool_parameters(schema, parameter_format=parameter_format) + + assert effective_format == "json" + assert parameters == original + parameters["$defs"]["Value"]["enum"].append("three") + assert schema == original + + +def test_format_tool_parameters_json_preserves_simple_schema(): + schema = {"type": "object", "properties": {"value": {"type": "string", "title": "Value"}}} + + assert _format_tool_parameters(schema, parameter_format="json") == ("json", schema) + + +@pytest.mark.parametrize("schema", [{}, {"properties": {"value": {"type": "string"}}}]) +def test_format_tool_parameters_does_not_treat_unconstrained_schema_as_empty(schema): + assert _format_tool_parameters(schema, parameter_format="compact") == ("json", schema) + + +@pytest.mark.parametrize("parameter_format", ["other", "", None]) +def test_format_tool_parameters_rejects_unknown_format(parameter_format): + with pytest.raises(ValueError, match="parameter_format"): + _format_tool_parameters({"type": "object", "properties": {}}, parameter_format=parameter_format) + + def test_tool_decorator(): """Test the tool decorator.""" diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index bf5487ce0f8..fa96ee9d60f 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -26,6 +26,7 @@ Message, ResponseStream, WorkflowBuilder, + WorkflowCheckpoint, WorkflowCheckpointException, WorkflowContext, WorkflowConvergenceException, @@ -1769,3 +1770,108 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None # endregion + + +# --------------------------------------------------------------------------- +# Pause checkpoint resolution (AG-UI interrupt metadata) +# --------------------------------------------------------------------------- + +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"}, + ) + 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, response + await ctx.yield_output("done") # type: ignore[arg-type] # pyrefly: ignore[bad-argument-type] # ty: ignore[invalid-argument-type] + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_prefers_runner_over_shared_latest() -> None: + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + assert pause_id is not None + + competing = WorkflowCheckpoint( + workflow_name=workflow.name, + graph_signature_hash="competing", + pending_request_info_events={}, + timestamp="9999-01-01T00:00:00+00:00", + ) + await storage.save(competing) + + # Workflow owns the run baseline captured at run() start. + resolved = await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + ) + assert resolved == pause_id + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_requires_run_scoped_change_without_storage() -> None: + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor()).build() + workflow._runner._previous_checkpoint_id = "leftover" # pyright: ignore[reportPrivateUsage] + workflow._run_baseline_checkpoint_id = "leftover" # pyright: ignore[reportPrivateUsage] + + assert await workflow.resolve_pause_checkpoint_id({"approval-1"}) is None + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + baseline_checkpoint_id=None, + ) + == "leftover" + ) + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_excludes_restored_checkpoint() -> None: + """After resume, the restored id must not be re-advertised if a new pause save fails.""" + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + + async for _ in workflow.run("go", stream=True): + pass + restored = workflow.get_last_checkpoint_id() + assert restored is not None + + # Simulate a resume that restored `restored` but did not persist a newer pause. + workflow._run_baseline_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + workflow._restored_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + workflow._runner._previous_checkpoint_id = restored # pyright: ignore[reportPrivateUsage] + + assert ( + await workflow.resolve_pause_checkpoint_id( + {"approval-1"}, + checkpoint_storage=storage, + known_checkpoint_id=restored, + ) + is None + ) + + +@pytest.mark.asyncio +async def test_resolve_pause_checkpoint_id_uses_builder_storage() -> None: + storage = InMemoryCheckpointStorage() + workflow = WorkflowBuilder(start_executor=_ApprovalExecutor(), checkpoint_storage=storage).build() + async for _ in workflow.run("go", stream=True): + pass + + pause_id = workflow.get_last_checkpoint_id() + resolved = await workflow.resolve_pause_checkpoint_id({"approval-1"}) + assert resolved == pause_id diff --git a/python/packages/hyperlight/README.md b/python/packages/hyperlight/README.md index 02332e100f4..a27ac0248a8 100644 --- a/python/packages/hyperlight/README.md +++ b/python/packages/hyperlight/README.md @@ -118,6 +118,40 @@ codeact = HyperlightCodeActProvider( ) ``` +### Sandbox tool parameter descriptions + +Both `HyperlightExecuteCodeTool` and `HyperlightCodeActProvider` accept the +keyword-only `tool_description_format` option. The default, `"compact"`, includes +each parameter's scalar type, required/optional status, description, enum values, +and default when present. Use `"json"` to include the complete JSON Schema: + +```python +execute_code = HyperlightExecuteCodeTool( + tools=[compute], + tool_description_format="json", +) + +codeact = HyperlightCodeActProvider( + tools=[compute], + tool_description_format={"compute": "json", "send_email": "compact"}, +) +``` + +A string applies to every registered tool. A mapping selects formats by exact, +case-sensitive tool name; missing names use `"compact"`. Mappings are copied at +construction and when creating run-scoped tools, and entries for unregistered +tools are retained for later registration. + +Compact mode automatically falls back to full JSON Schema, with an explanatory +note, when a schema cannot be represented faithfully (for example, nested objects, +arrays, references, or additional constraints). No schema details are discarded. +Only `"compact"` and `"json"` are accepted; `None` is not supported. + +This option affects `HyperlightExecuteCodeTool.description`, or the injected +run tool's `.description` when using `HyperlightCodeActProvider`. It does not +change the short CodeAct instructions, the `execute_code` input schema, sandbox +execution, or runtime caching. + ### Output attachment limits Files written under `/output` are returned as inline data attachments. Hyperlight diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py b/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py index 3618cc9e6b8..55ec0254df0 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_execute_code_tool.py @@ -10,14 +10,14 @@ import stat import threading import time -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor from contextlib import suppress from copy import copy from dataclasses import dataclass from pathlib import Path, PurePosixPath from tempfile import TemporaryDirectory -from typing import Any, Protocol, TypeGuard, TypeVar, cast +from typing import Any, Literal, Protocol, TypeGuard, TypeVar, cast from urllib.parse import urlparse from agent_framework import Content, FunctionTool @@ -1505,13 +1505,43 @@ def _build_sandbox() -> tuple[Any, Any]: ) +def _normalize_tool_description_format( + value: object, +) -> Literal["compact", "json"] | dict[str, Literal["compact", "json"]]: + if isinstance(value, str): + if value not in ("compact", "json"): + raise ValueError("tool_description_format must be 'compact' or 'json'.") + return value + if not isinstance(value, Mapping): + raise TypeError("tool_description_format must be 'compact', 'json', or a mapping of tool names to formats.") + + result: dict[str, Literal["compact", "json"]] = {} + for name, parameter_format in cast(Mapping[object, object], value).items(): + if not isinstance(name, str): + raise TypeError("tool_description_format mapping keys must be tool name strings.") + if not isinstance(parameter_format, str): + raise TypeError(f"tool_description_format for {name!r} must be a string: 'compact' or 'json'.") + if parameter_format not in ("compact", "json"): + raise ValueError(f"tool_description_format for {name!r} must be 'compact' or 'json'.") + result[name] = parameter_format + return result + + class HyperlightExecuteCodeTool(FunctionTool): - """Execute Python code inside a Hyperlight sandbox.""" + """Execute Python code inside a Hyperlight sandbox. + + Keyword Args: + tool_description_format: Parameter documentation in ``.description``: ``"compact"`` (default) + or ``"json"``, globally or mapped by exact, case-sensitive tool name. Missing names use + compact format. Schemas that cannot be represented faithfully in compact form use JSON Schema. + Mappings are copied, including entries for tools registered later. + """ def __init__( self, *, tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None, + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", approval_mode: ApprovalMode | None = None, workspace_root: str | Path | None = None, file_mounts: FileMountInput | Sequence[FileMountInput] | None = None, @@ -1524,6 +1554,7 @@ def __init__( module_path: str | None = None, _registry: SandboxRuntime | None = None, ) -> None: + normalized_description_format = _normalize_tool_description_format(tool_description_format) max_output_files = _validate_positive_integer(name="max_output_files", value=max_output_files) max_output_file_bytes = _validate_positive_integer(name="max_output_file_bytes", value=max_output_file_bytes) max_output_total_bytes = _validate_positive_integer(name="max_output_total_bytes", value=max_output_total_bytes) @@ -1535,6 +1566,9 @@ def __init__( input_model=EXECUTE_CODE_INPUT_SCHEMA, ) self._state_lock = threading.RLock() + self._tool_description_format: Literal["compact", "json"] | dict[str, Literal["compact", "json"]] = ( + normalized_description_format + ) self._registry = _registry or _SandboxRegistry() self._default_approval_mode: ApprovalMode = approval_mode or "never_require" self._workspace_root = _resolve_workspace_root(workspace_root) @@ -1571,6 +1605,7 @@ def description(self) -> str: workspace_enabled=self._workspace_root is not None, mounted_paths=[_display_mount_path(mount.mount_path) for mount in self._file_mounts.values()], allowed_domains=allowed_domains, + tool_description_format=self._tool_description_format, ) @description.setter @@ -1692,6 +1727,7 @@ def create_run_tool(self) -> HyperlightExecuteCodeTool: return HyperlightExecuteCodeTool( tools=self.get_tools(), + tool_description_format=self._tool_description_format, approval_mode=self._default_approval_mode, workspace_root=self._workspace_root, file_mounts=file_mounts or None, @@ -1717,6 +1753,11 @@ def build_serializable_state(self) -> dict[str, Any]: "max_output_file_bytes": config.max_output_file_bytes, "max_output_total_bytes": config.max_output_total_bytes, "tool_names": [tool_obj.name for tool_obj in config.tools], + "tool_description_format": ( + dict(self._tool_description_format) + if isinstance(self._tool_description_format, dict) + else self._tool_description_format + ), "filesystem_enabled": config.filesystem_enabled, "workspace_root": str(config.workspace_root) if config.workspace_root is not None else None, "file_mounts": [ diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_instructions.py b/python/packages/hyperlight/agent_framework_hyperlight/_instructions.py index c44a1830626..919afc6fc30 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_instructions.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_instructions.py @@ -2,24 +2,53 @@ from __future__ import annotations -from collections.abc import Sequence +import json +from collections.abc import Mapping, Sequence +from typing import Literal from agent_framework import FunctionTool +from agent_framework._tools import _format_tool_parameters # pyright: ignore[reportPrivateUsage] from ._types import AllowedDomain -def _format_tool_summaries(tools: Sequence[FunctionTool]) -> str: +def _format_tool_summaries( + tools: Sequence[FunctionTool], + *, + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", +) -> str: if not tools: return "- No tools are currently registered inside the sandbox." lines: list[str] = [] for tool_obj in tools: - parameters = tool_obj.parameters().get("properties", {}) - parameter_names = [name for name in parameters if isinstance(name, str)] - parameter_summary = ", ".join(parameter_names) if parameter_names else "none" + requested_format = ( + tool_description_format + if isinstance(tool_description_format, str) + else tool_description_format.get(tool_obj.name, "compact") + ) + effective_format, parameters = _format_tool_parameters(tool_obj.parameters(), parameter_format=requested_format) description = str(tool_obj.description or "").strip() or "No description provided." - lines.append(f"- `{tool_obj.name}`: {description} Parameters: {parameter_summary}.") + lines.append(f"- `{tool_obj.name}`: {description}") + if effective_format == "json": + if requested_format == "compact": + lines.append( + " Using JSON Schema because the parameter schema cannot be represented faithfully in compact form." + ) + lines.extend([" Parameters (JSON Schema):", "```json", json.dumps(parameters, indent=2), "```"]) + elif not parameters: + lines.append(" Parameters: none.") + else: + for name, parameter in parameters.items(): + requirement = "required" if parameter["required"] else "optional" + line = f" - `{name}` ({parameter['type']}, {requirement})" + if parameter.get("description"): + line += f": {parameter['description']}" + if "enum" in parameter: + line += f" Allowed values: {json.dumps(parameter['enum'])}." + if "default" in parameter: + line += f" Default: {json.dumps(parameter['default'])}." + lines.append(line) return "\n".join(lines) @@ -107,6 +136,7 @@ def build_execute_code_description( workspace_enabled: bool, mounted_paths: Sequence[str], allowed_domains: Sequence[AllowedDomain], + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", ) -> str: """Build the dynamic execute_code tool description for standalone usage.""" filesystem_text = _format_filesystem_capabilities( @@ -126,7 +156,7 @@ def build_execute_code_description( tool name. Registered sandbox tools: -{_format_tool_summaries(tools)} +{_format_tool_summaries(tools, tool_description_format=tool_description_format)} Filesystem capabilities: {filesystem_text} diff --git a/python/packages/hyperlight/agent_framework_hyperlight/_provider.py b/python/packages/hyperlight/agent_framework_hyperlight/_provider.py index 7d26f295ef8..63394115299 100644 --- a/python/packages/hyperlight/agent_framework_hyperlight/_provider.py +++ b/python/packages/hyperlight/agent_framework_hyperlight/_provider.py @@ -2,9 +2,9 @@ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from pathlib import Path -from typing import Any +from typing import Any, Literal from agent_framework import AgentSession, ContextProvider, FunctionTool, SessionContext from agent_framework._telemetry import mark_feature_used @@ -22,7 +22,14 @@ class HyperlightCodeActProvider(ContextProvider): - """Inject a Hyperlight-backed CodeAct surface using provider-owned tools.""" + """Inject a Hyperlight-backed CodeAct surface using provider-owned tools. + + Keyword Args: + tool_description_format: Parameter documentation in the injected tool's ``.description``: + ``"compact"`` (default) or ``"json"``, globally or mapped by exact, case-sensitive tool name. + Missing names use compact format; rich schemas fall back to full JSON Schema. + Mappings are copied, including entries for tools registered later. + """ DEFAULT_SOURCE_ID = "hyperlight_codeact" @@ -31,6 +38,7 @@ def __init__( source_id: str = DEFAULT_SOURCE_ID, *, tools: FunctionTool | Callable[..., Any] | Sequence[FunctionTool | Callable[..., Any]] | None = None, + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", approval_mode: ApprovalMode | None = None, workspace_root: str | Path | None = None, file_mounts: FileMountInput | Sequence[FileMountInput] | None = None, @@ -46,6 +54,7 @@ def __init__( super().__init__(source_id) self._execute_code_tool = HyperlightExecuteCodeTool( tools=tools, + tool_description_format=tool_description_format, approval_mode=approval_mode, workspace_root=workspace_root, file_mounts=file_mounts, diff --git a/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py b/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py index c9abad2af7b..b9ed18feceb 100644 --- a/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py +++ b/python/packages/hyperlight/tests/hyperlight/test_hyperlight_codeact.py @@ -17,10 +17,12 @@ import threading import time from collections.abc import Awaitable, Callable, Coroutine, Generator, Mapping, Sequence +from copy import deepcopy from dataclasses import dataclass from pathlib import Path from tempfile import TemporaryDirectory -from typing import Any, cast +from types import MappingProxyType +from typing import Any, Literal, cast from unittest.mock import patch import pytest @@ -1880,6 +1882,222 @@ def test_execute_code_tool_description_contains_call_tool_guidance(tmp_path: Pat assert "github.com" in description +@pytest.fixture +def documented_tool() -> FunctionTool: + return FunctionTool( + name="lookup", + description="Look up an item.", + func=lambda **kwargs: kwargs, + input_model={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search text"}, + "limit": {"type": "integer", "description": "Maximum results", "default": 10}, + "order": {"type": "string", "enum": ["ascending", "descending"], "default": "ascending"}, + }, + "required": ["query"], + }, + ) + + +@pytest.mark.parametrize( + "mode", + ["compact", {}, {"lookup": "compact"}, {"LOOKUP": "json"}, MappingProxyType({"inactive": "json"})], +) +def test_description_format_compact_preserves_parameter_metadata(documented_tool: FunctionTool, mode: Any) -> None: + execute_code = HyperlightExecuteCodeTool( + tools=[documented_tool], tool_description_format=mode, _registry=_FakeRuntime() + ) + description = execute_code.description + + assert "- `lookup`: Look up an item." in description + assert "- `query` (string, required): Search text" in description + assert "- `limit` (integer, optional): Maximum results Default: 10." in description + expected_order = '- `order` (string, optional) Allowed values: ["ascending", "descending"]. Default: "ascending".' + assert expected_order in description + assert "Parameters (JSON Schema)" not in description + assert execute_code.to_dict()["description"] == description + assert execute_code.parameters() == execute_code_module.EXECUTE_CODE_INPUT_SCHEMA + + +@pytest.mark.parametrize("mode", ["json", {"lookup": "json"}]) +@pytest.mark.parametrize("rich_schema", [False, True]) +def test_description_format_json_preserves_full_schema( + documented_tool: FunctionTool, mode: Any, rich_schema: bool +) -> None: + schema = deepcopy(documented_tool.parameters()) + if rich_schema: + schema["additionalProperties"] = False + schema["properties"]["filters"] = {"type": "array", "items": {"type": "integer", "minimum": 1}} + registered_tool = FunctionTool(name="lookup", description="", func=lambda **kwargs: kwargs, input_model=schema) + original_schema = deepcopy(registered_tool.parameters()) + execute_code = HyperlightExecuteCodeTool( + tools=[registered_tool], tool_description_format=mode, _registry=_FakeRuntime() + ) + + description = execute_code.description + rendered_schema = json.loads(description.split("```json\n", 1)[1].split("\n```", 1)[0]) + assert rendered_schema == original_schema + assert "Parameters (JSON Schema)" in description + assert "cannot be represented faithfully" not in description + assert registered_tool.parameters() == original_schema + assert execute_code.to_dict()["description"] == description + + +def test_description_format_rich_schema_falls_back_without_mutation() -> None: + schema = { + "type": "object", + "properties": { + "record": { + "type": "object", + "properties": {"ids": {"type": "array", "items": {"type": "integer"}}}, + "required": ["ids"], + } + }, + "required": ["record"], + "additionalProperties": False, + } + original_schema = deepcopy(schema) + registered_tool = FunctionTool(name="nested", description="", func=lambda **kwargs: kwargs, input_model=schema) + execute_code = HyperlightExecuteCodeTool(tools=[registered_tool], _registry=_FakeRuntime()) + + description = execute_code.description + assert ( + "Using JSON Schema because the parameter schema cannot be represented faithfully in compact form." + in description + ) + rendered_schema = json.loads(description.split("```json\n", 1)[1].split("\n```", 1)[0]) + assert rendered_schema == registered_tool.parameters() == original_schema + assert schema == original_schema + assert execute_code.to_dict()["description"] == description + assert execute_code.build_serializable_state()["tool_description_format"] == "compact" + + +def test_description_format_mixed_mapping_and_dynamic_registration(documented_tool: FunctionTool) -> None: + modes: dict[str, Literal["compact", "json"]] = {"lookup": "json", "compute": "compact", "future": "json"} + execute_code = HyperlightExecuteCodeTool(tools=[compute], tool_description_format=modes, _registry=_FakeRuntime()) + run_tool = execute_code.create_run_tool() + assert execute_code._tool_description_format is not modes + assert run_tool._tool_description_format is not execute_code._tool_description_format + modes["lookup"] = "compact" + execute_code.add_tools([documented_tool, dangerous_compute]) + + description = execute_code.description + assert description.count("Parameters (JSON Schema)") == 1 + assert "- `a` (integer, required)" in description + assert "- `dangerous_compute`: No description provided." in description + assert "lookup" not in run_tool.description + run_tool.add_tools(documented_tool) + assert "Parameters (JSON Schema)" in run_tool.description + execute_code.remove_tool("lookup") + execute_code.add_tools(documented_tool) + assert "Parameters (JSON Schema)" in execute_code.description + state = execute_code.build_serializable_state() + assert state["tool_description_format"] == {"lookup": "json", "compute": "compact", "future": "json"} + state["tool_description_format"]["lookup"] = "compact" + assert "Parameters (JSON Schema)" in execute_code.description + assert json.loads(json.dumps(run_tool.build_serializable_state()))["tool_description_format"]["future"] == "json" + + +@pytest.mark.parametrize("mode", ["compact", "json", {"compute": "json", "future": "compact"}]) +async def test_description_format_provider_run_state_and_snapshot(mode: Any) -> None: + provider = HyperlightCodeActProvider(tools=[compute], tool_description_format=mode, _registry=_FakeRuntime()) + expected_mode = deepcopy(mode) + if isinstance(mode, dict): + mode["compute"] = "compact" + context = _FakeSessionContext() + state: dict[str, Any] = {} + + await provider.before_run(agent=object(), session=None, context=cast(Any, context), state=state) + run_tool = context.tools[0][1][0] + assert isinstance(run_tool, HyperlightExecuteCodeTool) + assert state[provider.source_id]["tool_description_format"] == expected_mode + assert run_tool.build_serializable_state() == state[provider.source_id] + assert run_tool.to_dict()["description"] == run_tool.description + assert ("Parameters (JSON Schema)" in run_tool.description) == (expected_mode != "compact") + assert "Parameters (JSON Schema)" not in context.instructions[0][1] + assert "compute" not in context.instructions[0][1] + assert run_tool._tool_description_format == provider._execute_code_tool._tool_description_format + if isinstance(expected_mode, dict): + assert run_tool._tool_description_format is not provider._execute_code_tool._tool_description_format + provider.clear_tools() + assert "- `compute`:" in run_tool.description + assert "No tools are currently registered" in provider._execute_code_tool.description + json.dumps(state) + + +@pytest.mark.parametrize( + ("mode", "error_type"), + [ + ("verbose", ValueError), + ("JSON", ValueError), + (" compact", ValueError), + ({"inactive": "verbose"}, ValueError), + (None, TypeError), + (1, TypeError), + (["compact"], TypeError), + ({1: "json"}, TypeError), + ({"compute": None}, TypeError), + ({"inactive": 1}, TypeError), + ], +) +@pytest.mark.parametrize("entry_point", [HyperlightExecuteCodeTool, HyperlightCodeActProvider]) +def test_description_format_rejects_invalid_inputs(mode: Any, error_type: type[Exception], entry_point: Any) -> None: + with pytest.raises(error_type, match="tool_description_format"): + entry_point(tool_description_format=mode, _registry=_FakeRuntime()) + + +def test_description_format_defaults_no_tools_and_zero_parameters() -> None: + execute_code = HyperlightExecuteCodeTool(_registry=_FakeRuntime()) + assert "- No tools are currently registered inside the sandbox." in execute_code.description + assert execute_code.build_serializable_state()["tool_description_format"] == "compact" + + execute_code.add_tools( + FunctionTool(name="noop", description="", func=lambda: None, input_model={"type": "object", "properties": {}}) + ) + assert "- `noop`: No description provided.\n Parameters: none." in execute_code.description + assert "call_tool(name, **kwargs)" in execute_code.description + assert "arguments only. Do not pass a dict or any other positional arguments" in execute_code.description + assert "Filesystem access is unavailable" in execute_code.description + assert "Outbound network access is unavailable" in execute_code.description + + +async def test_description_format_does_not_change_runtime_config_or_input_schema() -> None: + runtime = _FakeRuntime() + compact_tool = HyperlightExecuteCodeTool(tools=[compute], _registry=runtime) + json_tool = HyperlightExecuteCodeTool(tools=[compute], tool_description_format="json", _registry=runtime) + for execute_code in (compact_tool, json_tool): + result = await execute_code.invoke(arguments={"code": "print(42)"}) + assert result[0].text == "ok" + assert execute_code.parameters() == execute_code_module.EXECUTE_CODE_INPUT_SCHEMA + compact_config, compact_code = runtime.calls[0] + json_config, json_code = runtime.calls[1] + assert compact_config == json_config + assert compact_config.cache_key() == json_config.cache_key() + assert not hasattr(compact_config, "tool_description_format") + assert compact_code == json_code == "print(42)" + assert compact_tool.build_instructions(tools_visible_to_model=False) == json_tool.build_instructions( + tools_visible_to_model=False + ) + + +async def test_description_format_reuses_sandbox_cache(monkeypatch: pytest.MonkeyPatch) -> None: + _FakeSandbox.instances.clear() + monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandbox) + registry = execute_code_module._SandboxRegistry() + compact_tool = HyperlightExecuteCodeTool(tools=[compute], _registry=registry) + json_tool = HyperlightExecuteCodeTool(tools=[compute], tool_description_format="json", _registry=registry) + + try: + for execute_code in (compact_tool, json_tool): + await execute_code.invoke(arguments={"code": "None"}) + assert len(_FakeSandbox.instances) == 1 + assert set(_FakeSandbox.instances[0].registered_tools) == {"compute"} + assert _FakeSandbox.instances[0].restore_calls == ["snapshot", "snapshot"] + finally: + registry.close() + + async def test_execute_code_tool_executes_with_structured_content(monkeypatch: pytest.MonkeyPatch) -> None: _FakeSandbox.instances.clear() monkeypatch.setattr(execute_code_module, "_load_sandbox_class", lambda: _FakeSandbox) diff --git a/python/packages/monty/AGENTS.md b/python/packages/monty/AGENTS.md index 67c3f45d8cb..e366b78750b 100644 --- a/python/packages/monty/AGENTS.md +++ b/python/packages/monty/AGENTS.md @@ -36,6 +36,19 @@ from agent_framework.monty import ( - `file_mounts` — sequence of `FileMountInput` (str shorthand, `(host_path, mount_path)` tuple, or `FileMount`) - `resource_limits` — Monty `ResourceLimits` TypedDict +- `tool_description_format` — keyword-only + `Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact"`. + Selects parameter documentation in both the description and instructions. + Compact includes scalar types, required/optional status, descriptions, enums, + and defaults; schemas that cannot be represented faithfully fall back to full + JSON Schema with an explanatory note. `"json"` always includes the full schema. + Mappings use exact, case-sensitive tool names, with compact for missing names; + mappings are copied on construction and per-run snapshots, retaining inactive + names for dynamic registration. Unsupported choices raise `ValueError`; invalid + input types (including `None`), non-string keys, and non-string choices raise + `TypeError`. Serializable state stores the configured string or mapping as + `tool_description_format`. This affects documentation only, not type checking + or invocation. Tool-management methods on both classes: `add_tools`, `get_tools`, `remove_tool`, `clear_tools`. Mount-management methods: `add_file_mounts`, diff --git a/python/packages/monty/README.md b/python/packages/monty/README.md index d6f5a8475a4..d1c2bc3e606 100644 --- a/python/packages/monty/README.md +++ b/python/packages/monty/README.md @@ -105,6 +105,32 @@ agent = Agent( ) ``` +### Tool parameter documentation + +Both `MontyCodeActProvider` and `MontyExecuteCodeTool` accept the keyword-only +`tool_description_format` parameter. It controls the registered tools' parameter +documentation in both the `execute_code` description and CodeAct instructions: + +- `"compact"` (default) lists scalar parameter types, required/optional status, + descriptions, enum choices, and defaults. +- `"json"` includes the complete parameter JSON Schema. +- A mapping such as `{"compute": "json", "send_email": "compact"}` selects formats + by exact, case-sensitive tool name. Names missing from the mapping use compact. + +Compact automatically falls back to complete JSON Schema, with an explanatory +note, for schemas it cannot represent faithfully, such as nested objects, arrays, +references, unions, or additional constraints. Parameter schemas and runtime +type checking are not changed. + +Mappings are copied at construction and for each run snapshot. Entries for +currently unregistered tools are retained for later `add_tools` calls, including +after removal or clearing of the registry. State snapshots include the configured +string or mapping in `tool_description_format`. + +Only `"compact"`, `"json"`, or mappings from string names to these values are +accepted: unsupported choices raise `ValueError`; `None`, invalid input types, +non-string mapping keys, and non-string choices raise `TypeError`. + ### Host tool lifetime Registered `FunctionTool` instances retain their invocation and exception counters diff --git a/python/packages/monty/agent_framework_monty/_execute_code_tool.py b/python/packages/monty/agent_framework_monty/_execute_code_tool.py index 18e0ab85bd0..a818dfb9246 100644 --- a/python/packages/monty/agent_framework_monty/_execute_code_tool.py +++ b/python/packages/monty/agent_framework_monty/_execute_code_tool.py @@ -14,10 +14,10 @@ import json import mimetypes -from collections.abc import Callable, Iterator, Sequence +from collections.abc import Callable, Iterator, Mapping, Sequence from functools import partial from pathlib import Path, PurePosixPath -from typing import Any, cast +from typing import Any, Literal, cast from agent_framework import Content, FunctionTool from agent_framework._tools import ApprovalMode, normalize_tools @@ -68,6 +68,27 @@ def _collect_tools(*tool_groups: Any) -> list[FunctionTool]: return list(tools_by_name.values()) +def _normalize_tool_description_format( + value: object, +) -> Literal["compact", "json"] | dict[str, Literal["compact", "json"]]: + if isinstance(value, str): + if value not in ("compact", "json"): + raise ValueError("tool_description_format must be 'compact', 'json', or a mapping of tool names to these.") + return value + if not isinstance(value, Mapping): + raise TypeError("tool_description_format must be a string or a mapping of tool names to 'compact' or 'json'.") + normalized: dict[str, Literal["compact", "json"]] = {} + for name, choice in cast("Mapping[object, object]", value).items(): + if not isinstance(name, str): + raise TypeError("tool_description_format mapping keys must be strings.") + if not isinstance(choice, str): + raise TypeError(f"tool_description_format[{name!r}] must be a string ('compact' or 'json').") + if choice not in ("compact", "json"): + raise ValueError(f"tool_description_format[{name!r}] must be 'compact' or 'json'; got {choice!r}.") + normalized[name] = choice + return normalized + + def _resolve_execute_code_approval_mode( *, base_approval_mode: ApprovalMode, @@ -186,6 +207,12 @@ class MontyExecuteCodeTool(FunctionTool): ``resource_limits`` is forwarded to Monty's ``ResourceLimits`` to cap CPU time, memory, output size, recursion depth, and GC frequency. + ``tool_description_format`` controls parameter documentation in both the + description and instructions: ``"compact"`` (default) or ``"json"`` globally, + or a mapping of exact, case-sensitive tool names to either format. Missing + names use compact; rich schemas fall back to full JSON Schema. Mappings + are copied and retain entries for tools registered later. + All mutators (``add_tools``, ``add_file_mounts`` etc.) must be called from the same task/thread that owns the tool. Monty itself runs on the event loop, so no internal locking is needed. @@ -199,6 +226,7 @@ def __init__( workspace_root: str | Path | None = None, file_mounts: FileMountInput | Sequence[FileMountInput] | None = None, resource_limits: dict[str, Any] | None = None, + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", ) -> None: super().__init__( name=EXECUTE_CODE_TOOL_NAME, @@ -208,6 +236,9 @@ def __init__( input_model=EXECUTE_CODE_INPUT_SCHEMA, ) self._default_approval_mode: ApprovalMode = approval_mode or "never_require" + self._tool_description_format: Literal["compact", "json"] | dict[str, Literal["compact", "json"]] = ( + _normalize_tool_description_format(tool_description_format) + ) self._managed_tools: list[FunctionTool] = [] self._workspace_root: Path | None = ( _resolve_existing_directory(workspace_root) if workspace_root is not None else None @@ -230,6 +261,7 @@ def description(self) -> str: return build_execute_code_description( tools=self._managed_tools, mounts=self._effective_mounts(), + tool_description_format=self._tool_description_format, ) @description.setter @@ -307,6 +339,7 @@ def build_instructions(self, *, tools_visible_to_model: bool) -> str: tools=list(self._managed_tools), tools_visible_to_model=tools_visible_to_model, mounts=self._effective_mounts(), + tool_description_format=self._tool_description_format, ) def create_run_tool(self) -> MontyExecuteCodeTool: @@ -317,6 +350,7 @@ def create_run_tool(self) -> MontyExecuteCodeTool: workspace_root=self._workspace_root, file_mounts=list(self._file_mounts.values()) or None, resource_limits=self._resource_limits, + tool_description_format=self._tool_description_format, ) def build_serializable_state(self) -> dict[str, Any]: @@ -330,6 +364,11 @@ def build_serializable_state(self) -> dict[str, Any]: "runtime": "monty", "approval_mode": approval_mode, "tool_names": [tool_obj.name for tool_obj in self._managed_tools], + "tool_description_format": ( + dict(self._tool_description_format) + if isinstance(self._tool_description_format, dict) + else self._tool_description_format + ), "workspace_root": str(self._workspace_root) if self._workspace_root is not None else None, "file_mounts": [ { diff --git a/python/packages/monty/agent_framework_monty/_instructions.py b/python/packages/monty/agent_framework_monty/_instructions.py index c560e356d31..6503b556ad6 100644 --- a/python/packages/monty/agent_framework_monty/_instructions.py +++ b/python/packages/monty/agent_framework_monty/_instructions.py @@ -4,24 +4,53 @@ from __future__ import annotations -from collections.abc import Sequence +import json +from collections.abc import Mapping, Sequence +from typing import Literal from agent_framework import FunctionTool +from agent_framework._tools import _format_tool_parameters # pyright: ignore[reportPrivateUsage] from ._types import FileMount -def _format_tool_summaries(tools: Sequence[FunctionTool]) -> str: +def _format_tool_summaries( + tools: Sequence[FunctionTool], + *, + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", +) -> str: if not tools: return "- No tools are currently registered." lines: list[str] = [] for tool_obj in tools: - parameters = tool_obj.parameters().get("properties", {}) - parameter_names = [name for name in parameters if isinstance(name, str)] - parameter_summary = ", ".join(parameter_names) if parameter_names else "none" + requested_format = ( + tool_description_format + if isinstance(tool_description_format, str) + else tool_description_format.get(tool_obj.name, "compact") + ) + effective_format, parameters = _format_tool_parameters(tool_obj.parameters(), parameter_format=requested_format) description = str(tool_obj.description or "").strip() or "No description provided." - lines.append(f"- `{tool_obj.name}`: {description} Parameters: {parameter_summary}.") + lines.append(f"- `{tool_obj.name}`: {description}") + if effective_format == "json": + if requested_format == "compact": + lines.append( + " Using JSON Schema because the parameter schema cannot be represented faithfully in compact form." + ) + lines.extend([" Parameters (JSON Schema):", " ```json", json.dumps(parameters, indent=2), " ```"]) + elif not parameters: + lines.append(" Parameters: none.") + else: + for name, parameter in parameters.items(): + required = "required" if parameter["required"] else "optional" + line = f" - `{name}` ({parameter['type']}, {required})" + if parameter.get("description"): + line += f": {parameter['description']}" + if "enum" in parameter: + line += f" Allowed values: {json.dumps(parameter['enum'])}." + if "default" in parameter: + line += f" Default: {json.dumps(parameter['default'])}." + lines.append(line) return "\n".join(lines) @@ -56,9 +85,10 @@ def build_codeact_instructions( tools: Sequence[FunctionTool], tools_visible_to_model: bool, mounts: Sequence[FileMount] = (), + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", ) -> str: """Build dynamic CodeAct instructions for the effective Monty tool set.""" - tool_summaries = _format_tool_summaries(tools) + tool_summaries = _format_tool_summaries(tools, tool_description_format=tool_description_format) filesystem_text = _format_filesystem_capabilities(mounts) usage_note = ( @@ -99,9 +129,10 @@ def build_execute_code_description( *, tools: Sequence[FunctionTool], mounts: Sequence[FileMount] = (), + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", ) -> str: """Build the dynamic ``execute_code`` tool description for standalone usage.""" - tool_summaries = _format_tool_summaries(tools) + tool_summaries = _format_tool_summaries(tools, tool_description_format=tool_description_format) filesystem_text = _format_filesystem_capabilities(mounts) return f"""Execute Python code in a Monty interpreter. diff --git a/python/packages/monty/agent_framework_monty/_provider.py b/python/packages/monty/agent_framework_monty/_provider.py index 6e45e958a19..51ababa5401 100644 --- a/python/packages/monty/agent_framework_monty/_provider.py +++ b/python/packages/monty/agent_framework_monty/_provider.py @@ -4,9 +4,9 @@ from __future__ import annotations -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from pathlib import Path -from typing import Any +from typing import Any, Literal from agent_framework import AgentSession, ContextProvider, FunctionTool, SessionContext from agent_framework._telemetry import mark_feature_used @@ -24,6 +24,13 @@ class MontyCodeActProvider(ContextProvider): the subset of capabilities that apply to the Monty interpreter: ``tools``, ``approval_mode``, ``workspace_root``, ``file_mounts``, and ``resource_limits`` (Monty-only). + + ``tool_description_format`` controls parameter documentation in both + injected instructions and the tool description. Pass ``"compact"`` (the + default) or ``"json"`` globally, or a mapping of exact, case-sensitive tool + names to either format. Missing names use compact; rich schemas fall back + to full JSON Schema. Mappings are copied on construction and for each run, + retaining entries for tools registered later. """ DEFAULT_SOURCE_ID = "monty_codeact" @@ -37,6 +44,7 @@ def __init__( workspace_root: str | Path | None = None, file_mounts: FileMountInput | Sequence[FileMountInput] | None = None, resource_limits: dict[str, Any] | None = None, + tool_description_format: Literal["compact", "json"] | Mapping[str, Literal["compact", "json"]] = "compact", ) -> None: super().__init__(source_id) self._execute_code_tool = MontyExecuteCodeTool( @@ -45,6 +53,7 @@ def __init__( workspace_root=workspace_root, file_mounts=file_mounts, resource_limits=resource_limits, + tool_description_format=tool_description_format, ) def add_tools( diff --git a/python/packages/monty/tests/monty/test_monty_codeact.py b/python/packages/monty/tests/monty/test_monty_codeact.py index 35548130ae4..db1fc8c6750 100644 --- a/python/packages/monty/tests/monty/test_monty_codeact.py +++ b/python/packages/monty/tests/monty/test_monty_codeact.py @@ -11,12 +11,13 @@ import asyncio import json +import re import sys import types -from collections.abc import Iterable, Iterator +from collections.abc import Iterable, Iterator, Mapping from dataclasses import dataclass, field from pathlib import Path -from typing import Annotated, Any +from typing import Annotated, Any, Literal from unittest.mock import MagicMock import pytest @@ -307,6 +308,194 @@ def test_dynamic_description_reflects_registered_tools() -> None: assert "mul_tool" in description_updated +@pytest.fixture +def documented_tool() -> FunctionTool: + return FunctionTool( + name="documented", + description="Documented scalar parameters.", + func=lambda **kwargs: kwargs, + input_model={ + "type": "object", + "properties": { + "query": {"type": "string", "description": "Search text."}, + "count": {"type": "integer", "description": "Result count.", "default": 0}, + "enabled": {"type": "boolean", "default": False}, + "ratio": {"type": "number", "enum": [0.5, 1.0], "default": 0.5}, + "empty": {"type": "null", "default": None}, + }, + "required": ["query"], + }, + ) + + +def _json_schemas(text: str) -> list[dict[str, Any]]: + return [json.loads(schema) for schema in re.findall(r"```json\n(.*?)\n\s*```", text, re.DOTALL)] + + +@pytest.mark.parametrize("tools_visible_to_model", [False, True]) +@pytest.mark.parametrize("tool_description_format", ["compact", {}, {"documented": "compact"}]) +def test_compact_parameter_documentation( + documented_tool: FunctionTool, + tools_visible_to_model: bool, + tool_description_format: Any, +) -> None: + monty_tool = MontyExecuteCodeTool(tools=[documented_tool], tool_description_format=tool_description_format) + instructions = monty_tool.build_instructions(tools_visible_to_model=tools_visible_to_model) + expected_lines = [ + "- `documented`: Documented scalar parameters.", + "- `query` (string, required): Search text.", + "- `count` (integer, optional): Result count. Default: 0.", + "- `enabled` (boolean, optional) Default: false.", + "- `ratio` (number, optional) Allowed values: [0.5, 1.0]. Default: 0.5.", + "- `empty` (null, optional) Default: null.", + ] + for text in (monty_tool.description, instructions): + for line in expected_lines: + assert line in text + assert "JSON Schema" not in text + assert "await tool_name(param=value)" in text + assert "await call_tool('name', **kwargs)" in text + assert "asyncio.gather" in text + assert "type-checked" in text + usage_note = ( + "Some tools may also appear directly" + if tools_visible_to_model + else "Provider-owned sandbox tools are not exposed separately" + ) + assert usage_note in instructions + assert monty_tool.to_dict()["description"] == monty_tool.description + assert MontyExecuteCodeTool(tools=[documented_tool]).description == monty_tool.description + + +@pytest.mark.parametrize("tools_visible_to_model", [False, True]) +@pytest.mark.parametrize( + ("tool_description_format", "json_tool_names"), + [ + ("json", ["documented", "add_tool"]), + ({"documented": "json", "add_tool": "compact"}, ["documented"]), + ({"documented": "json"}, ["documented"]), + ({"Documented": "json", "unused": "json"}, []), + ], +) +def test_parameter_documentation_formats( + documented_tool: FunctionTool, + tools_visible_to_model: bool, + tool_description_format: Any, + json_tool_names: list[str], +) -> None: + tools = [documented_tool, add_tool] + monty_tool = MontyExecuteCodeTool(tools=tools, tool_description_format=tool_description_format) + expected = [registered.parameters() for registered in tools if registered.name in json_tool_names] + for text in ( + monty_tool.description, + monty_tool.build_instructions(tools_visible_to_model=tools_visible_to_model), + ): + assert _json_schemas(text) == expected + assert text.count("Parameters (JSON Schema)") == len(expected) + assert "Using JSON Schema because" not in text + if "add_tool" not in json_tool_names: + assert "- `a` (integer, required): First addend" in text + if "documented" not in json_tool_names: + assert "- `query` (string, required): Search text." in text + assert monty_tool.to_dict()["description"] == monty_tool.description + state = monty_tool.build_serializable_state() + assert state["tool_description_format"] == tool_description_format + assert json.loads(json.dumps(state)) == state + + +@pytest.mark.parametrize("tool_description_format", ["compact", "json", {"rich": "compact"}, {"rich": "json"}]) +def test_rich_parameter_schema_is_preserved(tool_description_format: Any) -> None: + schema: dict[str, Any] = { + "type": "object", + "properties": { + "items": {"type": "array", "items": {"$ref": "#/$defs/Entry"}, "minItems": 1}, + "limit": {"type": "integer", "minimum": 1, "default": 2}, + }, + "$defs": {"Entry": {"type": "object", "properties": {"value": {"type": "string"}}}}, + "required": ["items"], + "additionalProperties": False, + } + rich = FunctionTool(name="rich", description="Rich schema.", func=lambda **kwargs: kwargs, input_model=schema) + monty_tool = MontyExecuteCodeTool(tools=[rich], tool_description_format=tool_description_format) + requested = tool_description_format if isinstance(tool_description_format, str) else tool_description_format["rich"] + for text in ( + monty_tool.description, + monty_tool.build_instructions(tools_visible_to_model=False), + monty_tool.build_instructions(tools_visible_to_model=True), + ): + assert _json_schemas(text) == [rich.parameters()] + assert rich.parameters() == schema + assert ( + "Using JSON Schema because the parameter schema cannot be represented faithfully in compact form." in text + ) == (requested == "compact") + assert monty_tool.to_dict()["description"] == monty_tool.description + + +@pytest.mark.parametrize("constructor", [MontyExecuteCodeTool, MontyCodeActProvider]) +@pytest.mark.parametrize( + ("value", "error"), + [ + ("unknown", ValueError), + ("JSON", ValueError), + (" compact", ValueError), + ({"inactive": "invalid"}, ValueError), + ({"inactive": "JSON"}, ValueError), + (None, TypeError), + (1, TypeError), + (False, TypeError), + (["compact"], TypeError), + ([("name", "json")], TypeError), + ({1: "json"}, TypeError), + ({"name": None}, TypeError), + ({"name": 1}, TypeError), + ({"name": ["json"]}, TypeError), + ], +) +def test_tool_description_format_validation(constructor: Any, value: Any, error: type[Exception]) -> None: + with pytest.raises(error, match="tool_description_format"): + constructor(tool_description_format=value) + + +def test_description_format_mapping_is_copied_and_retained_for_dynamic_tools() -> None: + formats: dict[str, Literal["compact", "json"]] = {"add_tool": "json", "mul_tool": "json"} + monty_tool = MontyExecuteCodeTool(tools=[add_tool], tool_description_format=types.MappingProxyType(formats)) + run_tool = monty_tool.create_run_tool() + snapshot_description = run_tool.description + snapshot_instructions = run_tool.build_instructions(tools_visible_to_model=False) + formats.clear() + state = monty_tool.build_serializable_state() + state["tool_description_format"].clear() + assert monty_tool.build_serializable_state()["tool_description_format"] == {"add_tool": "json", "mul_tool": "json"} + assert run_tool._tool_description_format is not monty_tool._tool_description_format + monty_tool.add_tools(mul_tool) + assert _json_schemas(monty_tool.description) == [add_tool.parameters(), mul_tool.parameters()] + monty_tool.remove_tool("add_tool") + assert "- `add_tool`:" not in monty_tool.description + monty_tool.clear_tools() + assert "No tools are currently registered." in monty_tool.description + monty_tool.add_tools(add_tool) + assert _json_schemas(monty_tool.description) == [add_tool.parameters()] + assert run_tool.description == snapshot_description + assert run_tool.build_instructions(tools_visible_to_model=False) == snapshot_instructions + + +@pytest.mark.parametrize("tool_description_format", ["compact", "json", {}]) +def test_empty_registry_and_zero_parameter_documentation(tool_description_format: Any) -> None: + monty_tool = MontyExecuteCodeTool(tool_description_format=tool_description_format) + assert "- No tools are currently registered." in monty_tool.description + assert "- No tools are currently registered." in monty_tool.build_instructions(tools_visible_to_model=False) + empty = FunctionTool( + name="empty", description="", func=lambda: None, input_model={"type": "object", "properties": {}} + ) + monty_tool.add_tools(empty) + for text in (monty_tool.description, monty_tool.build_instructions(tools_visible_to_model=True)): + assert "- `empty`: No description provided." in text + if tool_description_format == "json": + assert _json_schemas(text) == [empty.parameters()] + else: + assert "Parameters: none." in text + + def test_create_run_tool_snapshots_current_state() -> None: monty_tool = MontyExecuteCodeTool(tools=[add_tool], approval_mode="never_require") run_tool = monty_tool.create_run_tool() @@ -329,6 +518,7 @@ def test_build_serializable_state_matches_effective_config() -> None: assert state["workspace_root"] is None assert state["file_mounts"] == [] assert state["resource_limits"] is None + assert state["tool_description_format"] == "compact" def test_file_mounts_normalized_and_round_tripped(tmp_path: Path) -> None: @@ -726,6 +916,48 @@ async def test_provider_injects_execute_code_tool_and_instructions() -> None: assert context.tools[0] is not provider._execute_code_tool # type: ignore[attr-defined] +@pytest.mark.parametrize("tool_description_format", ["compact", "json", {"add_tool": "json"}]) +async def test_provider_parameter_documentation_and_run_state(tool_description_format: Any) -> None: + provider = MontyCodeActProvider(tools=[add_tool, mul_tool], tool_description_format=tool_description_format) + context = SessionContext(input_messages=[]) + state: dict[str, Any] = {} + await provider.before_run(agent=MagicMock(), session=None, context=context, state=state) + run_tool = context.tools[0] + assert isinstance(run_tool, MontyExecuteCodeTool) + instructions = "\n".join(context.instructions) + expected_registry = instructions_module._format_tool_summaries( + [add_tool, mul_tool], tool_description_format=tool_description_format + ) + assert expected_registry in run_tool.description + assert expected_registry in instructions + assert "Provider-owned sandbox tools are not exposed separately" in instructions + assert state["monty_codeact"]["tool_description_format"] == tool_description_format + assert run_tool.to_dict()["description"] == run_tool.description + + +async def test_provider_copies_format_mapping_and_snapshots_dynamic_registry() -> None: + formats: dict[str, Literal["compact", "json"]] = {"add_tool": "json", "mul_tool": "json"} + configuration: Mapping[str, Literal["compact", "json"]] = types.MappingProxyType(formats) + provider = MontyCodeActProvider(tools=[add_tool], tool_description_format=configuration) + formats["add_tool"] = "compact" + first = SessionContext(input_messages=[]) + await provider.before_run(agent=MagicMock(), session=None, context=first, state={}) + first_tool = first.tools[0] + assert isinstance(first_tool, MontyExecuteCodeTool) + first_description = first_tool.description + assert _json_schemas(first_description) == [add_tool.parameters()] + provider.clear_tools() + provider.add_tools([mul_tool]) + second = SessionContext(input_messages=[]) + state: dict[str, Any] = {} + await provider.before_run(agent=MagicMock(), session=None, context=second, state=state) + second_tool = second.tools[0] + assert isinstance(second_tool, MontyExecuteCodeTool) + assert _json_schemas(second_tool.description) == [mul_tool.parameters()] + assert state["monty_codeact"]["tool_description_format"] == {"add_tool": "json", "mul_tool": "json"} + assert first_tool.description == first_description + + def test_provider_delegates_tool_management_to_internal_tool() -> None: provider = MontyCodeActProvider() provider.add_tools([add_tool, mul_tool]) diff --git a/python/packages/purview/README.md b/python/packages/purview/README.md index 40a53b839a1..3b596a75c13 100644 --- a/python/packages/purview/README.md +++ b/python/packages/purview/README.md @@ -314,8 +314,46 @@ except (PurviewAuthenticationError, PurviewRateLimitError, PurviewRequestError, --- +## Security Considerations + +### Identity is a trusted input + +Purview evaluates DLP policy **for a specific user**. The identity this integration resolves therefore +decides *which* policy is applied, and it is resolved in this order: + +1. The `user_id` from the configured credential's token, when the credential resolves to a user. +2. The `user_id` argument passed to the processor. +3. `message.additional_properties["user_id"]`. +4. `message.author_name`, when it is a GUID. + +Only source 1 is verified. Sources 2–4 are supplied by the hosting application, so **a host must not +populate them from data that has crossed a trust boundary**. If an end user, an upstream service or a +model response can influence `additional_properties["user_id"]` or `author_name`, that party can +select a different user's DLP policy — typically one with weaker rules — and evade enforcement. Where +identity must come from a request, derive it from a validated token on the server, never from the +request body. Prefer a user-delegated credential (source 1) whenever possible. + +The `purview_app_location` in `PurviewSettings` is trusted in the same way: it selects which policy +locations apply and must be configured by the host, not by the caller. + +### Fail-closed behaviour + +Policy evaluation fails closed. If no user id can be resolved, or the tenant or app location cannot be +determined, the processor raises rather than letting content through unevaluated. Use +`ignore_exceptions` if you deliberately want availability over enforcement — but understand that it +disables enforcement for every error, not just transient ones. + +### What is evaluated + +Every content item on a message is submitted for evaluation, not just its text: binary/data content is +sent as Purview binary content, and function calls, function results and other structured content are +serialized to text. Only `usage` content is skipped, because it carries token counts rather than user +data. + +--- + ## Notes -- **User Identification**: When the configured credential resolves to a user token, that token's `user_id` is used for per-user policy scoping. For app-token credentials, provide a `user_id` per request (e.g. in `Message(..., additional_properties={"user_id": ""})`). If no user_id is provided or inferred, policy evaluation is skipped. +- **User Identification**: When the configured credential resolves to a user token, that token's `user_id` is used for per-user policy scoping. For app-token credentials, provide a `user_id` per request (e.g. in `Message(..., additional_properties={"user_id": ""})`). If no user_id can be provided or inferred, the request fails rather than proceeding unevaluated — see [Security Considerations](#security-considerations). - **Blocking Messages**: Can be customized via `blocked_prompt_message` and `blocked_response_message` in `PurviewSettings`. By default, they are "Prompt blocked by policy" and "Response blocked by policy" respectively. - **Streaming Responses**: Post-response policy evaluation presently applies only to non-streaming chat responses. - **Error Handling**: Use `ignore_exceptions` and `ignore_payment_required` settings for graceful degradation. When enabled, errors are logged but don't fail the request. diff --git a/python/packages/purview/agent_framework_purview/_models.py b/python/packages/purview/agent_framework_purview/_models.py index 40df9d78cfe..76e1b88acc0 100644 --- a/python/packages/purview/agent_framework_purview/_models.py +++ b/python/packages/purview/agent_framework_purview/_models.py @@ -153,6 +153,7 @@ def serialize_flag( class DlpAction(str, Enum): BLOCK_ACCESS = "blockAccess" + RESTRICT_ACCESS = "restrictAccess" OTHER = "other" @@ -611,7 +612,7 @@ class ContentToProcess(_AliasSerializable): def __init__( self, - content_entries: list[ProcessConversationMetadata | MutableMapping[str, Any]], + content_entry: ProcessConversationMetadata | MutableMapping[str, Any], activity_metadata: ActivityMetadata | MutableMapping[str, Any], device_metadata: DeviceMetadata | MutableMapping[str, Any], integrated_app_metadata: IntegratedAppMetadata | MutableMapping[str, Any], @@ -619,8 +620,6 @@ def __init__( **kwargs: Any, ) -> None: # Extract aliased values from kwargs - if "contentEntries" in kwargs: - content_entries = kwargs["contentEntries"] if "activityMetadata" in kwargs: activity_metadata = kwargs["activityMetadata"] if "deviceMetadata" in kwargs: @@ -631,10 +630,11 @@ def __init__( protected_app_metadata = kwargs["protectedAppMetadata"] # Convert nested objects - entries = [ - e if isinstance(e, ProcessConversationMetadata) else ProcessConversationMetadata(**e) - for e in content_entries - ] + entry = ( + content_entry + if isinstance(content_entry, ProcessConversationMetadata) + else ProcessConversationMetadata(**content_entry) + ) if isinstance(activity_metadata, MutableMapping): activity_metadata = ActivityMetadata(**activity_metadata) if isinstance(device_metadata, MutableMapping): @@ -646,7 +646,7 @@ def __init__( # Call parent without explicit params with aliases super().__init__(**kwargs) - self.content_entries = entries + self.content_entries = [entry] self.activity_metadata = activity_metadata self.device_metadata = device_metadata self.integrated_app_metadata = integrated_app_metadata diff --git a/python/packages/purview/agent_framework_purview/_processor.py b/python/packages/purview/agent_framework_purview/_processor.py index d6c24a078f2..919e97723e3 100644 --- a/python/packages/purview/agent_framework_purview/_processor.py +++ b/python/packages/purview/agent_framework_purview/_processor.py @@ -1,13 +1,18 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import base64 +import binascii +import json import logging +import re import time import uuid from collections.abc import Iterable, MutableMapping +from copy import copy from typing import Any -from agent_framework import Message +from agent_framework import Content, Message from ._cache import CacheProvider, InMemoryCacheProvider, create_protection_scopes_cache_key from ._client import PurviewClient @@ -16,6 +21,7 @@ Activity, ActivityMetadata, ContentActivitiesRequest, + ContentBase, ContentToProcess, DeviceMetadata, DlpAction, @@ -31,6 +37,7 @@ ProtectionScopesRequest, ProtectionScopesResponse, ProtectionScopeState, + PurviewBinaryContent, PurviewTextContent, RestrictionAction, translate_activity, @@ -51,6 +58,126 @@ def _is_valid_guid(value: str | None) -> bool: return False +# RFC 2397: data:[][;base64],, where may be followed by any number of +# ";parameter=value" segments (for example "data:text/plain;charset=utf-8;base64,..."). The optional +# parameters have to be matched explicitly, otherwise a parameterised media type fails to decode and +# its payload would be submitted as a base64 string that no classifier can read. +_DATA_URI_PATTERN = re.compile( + r"^data:(?P[^;,]*)(?:;[^;,]+)*?;base64,(?P.*)$", + re.DOTALL | re.IGNORECASE, +) + +# Content types that carry no user data and therefore have nothing for DLP to classify. +# Every other content type must reach Purview; see _map_content. +_NON_EVALUATED_CONTENT_TYPES = frozenset({"usage"}) + + +def _serialize_for_evaluation(value: Any) -> str: + """Render an arbitrary content value as text so Purview can classify it.""" + if value is None: + return "" + if isinstance(value, str): + return value + try: + return json.dumps(value, default=str, sort_keys=True) + except (TypeError, ValueError): + return str(value) + + +def _decode_data_uri(uri: str | None) -> bytes | None: + """Return the raw bytes behind a base64 data URI, or None if it is not one.""" + if not uri: + return None + match = _DATA_URI_PATTERN.match(uri) + if not match: + return None + try: + return base64.b64decode(match.group("base64_data"), validate=True) + except (binascii.Error, ValueError): + return None + + +def _map_content(content: Content) -> ContentBase | None: + """Map a single message content item onto the Purview content type that fits it. + + Returns None only for content that carries no user data at all. Anything else is + mapped to a real content entry: a message must never reach the model with part of + its payload unevaluated. + """ + content_type = content.type + + if content_type in _NON_EVALUATED_CONTENT_TYPES: + return None + + # Empty data is skipped because the Graph APIs do not support it. A request with + # empty data cannot violate content policies because it contains no content. + if content_type in ("text", "text_reasoning"): + if not content.text: + return None + return PurviewTextContent(data=content.text) + + if content_type == "data": + uri = getattr(content, "uri", None) + if not uri: + return None + raw_data = _decode_data_uri(uri) + if raw_data is not None: + if not raw_data: + return None + return PurviewBinaryContent(data=raw_data) + # Not a base64 data URI after all: evaluate the serialized form rather than drop it. + return PurviewTextContent(data=_serialize_for_evaluation(content.to_dict())) + + # Everything else - uri, function_call, function_result and any content type added after this + # code was written - is serialized whole. Serializing the entire item rather than picking out + # named fields keeps additional_properties, which is a data channel in its own right, under + # evaluation. Skipping a content type would be a policy bypass. + return PurviewTextContent(data=_serialize_for_evaluation(content.to_dict())) + + +def _map_message_contents(message: Message) -> list[ContentBase]: + """Map every non-empty content item of a message to Purview content entries.""" + return [purview_content for content in message.contents if (purview_content := _map_content(content))] + + +def _is_blocking_action(action_info: DlpActionInfo) -> bool: + """Whether a policy action means the content must not be released. + + ``restrictAccess`` is not blocking on its own: it carries a separate ``restrictionAction`` that + selects the enforcement mode, which may be audit, warn or allow as well as block. Only an explicit + block mode withholds the content. + """ + return action_info.action == DlpAction.BLOCK_ACCESS or action_info.restriction_action == RestrictionAction.BLOCK + + +def _blocking_actions(dlp_actions: list[DlpActionInfo]) -> list[DlpActionInfo | MutableMapping[str, Any]]: + """Filter policy actions down to the ones that block.""" + return [action_info for action_info in dlp_actions if _is_blocking_action(action_info)] + + +def _normalize_location_value(data_type: str, value: str) -> str: + """Normalize a policy location value for comparison, according to its location type. + + Application ids (GUIDs) and domain names are case-insensitive, so those fold whole. URL values are + not. Only the scheme and the host are case-insensitive; the userinfo, path, query and fragment are + all case-sensitive, so folding a URL whole would let a scope for ``contoso.com/public`` match a + request for ``contoso.com/Public``. Location types that are not recognised fold whole, which matches + more scopes rather than fewer. + """ + if data_type.split(".")[-1].casefold().endswith("url"): + scheme, separator, remainder = value.partition("://") + if not separator: + scheme, separator, remainder = "", "", value + # The authority ends at the first path, query or fragment delimiter, whichever comes first. + delimiter = re.search(r"[/?#]", remainder) + authority_end = delimiter.start() if delimiter else len(remainder) + authority, tail = remainder[:authority_end], remainder[authority_end:] + # Credentials are case-sensitive, so only the host half of the authority folds. + userinfo, at_sign, host = authority.rpartition("@") + return f"{scheme.casefold()}{separator}{userinfo}{at_sign}{host.casefold()}{tail}" + return value.casefold() + + class ScopedContentProcessor: """Combine protection scopes, process content, and content activities logic.""" @@ -91,7 +218,7 @@ async def process_messages( resp = await self._process_with_scopes(req) if resp.policy_actions: for act in resp.policy_actions: - if act.action == DlpAction.BLOCK_ACCESS or act.restriction_action == RestrictionAction.BLOCK: + if _is_blocking_action(act): should_block = True break if should_block: @@ -141,23 +268,21 @@ async def _map_messages( if not resolved_user_id and resolved_author_name: resolved_user_id = resolved_author_name - # Return empty results if user_id is empty + # Fail closed: without a resolvable user identity no policy can be evaluated, + # so the content must not be allowed through unevaluated. if not resolved_user_id or not _is_valid_guid(resolved_user_id): - return results, None + raise ValueError( + "No user id provided or inferred for Purview request. Please provide an Entra user id in each " + "message, pass a user id to the processor, or configure the credential to authenticate to an " + "Entra user." + ) for m in messages: message_id = m.message_id or str(uuid.uuid4()) - content = PurviewTextContent(data=m.text or "") correlation_id = (session_id or str(uuid.uuid4())) + "@AF" - meta = ProcessConversationMetadata( - identifier=message_id, - content=content, - name=f"Agent Framework Message {message_id}", - is_truncated=False, - correlation_id=correlation_id, - # This would be c# ticks equivalent and needs to fit inside c# long - sequence_number=time.time_ns() // 100 + 621355968000000000, - ) + # This would be c# ticks equivalent and needs to fit inside c# long + base_sequence_number = time.time_ns() // 100 + 621355968000000000 + mapped_contents = _map_message_contents(m) activity_meta = ActivityMetadata(activity=activity) purview_app_location = self._settings.get("purview_app_location") @@ -187,21 +312,30 @@ async def _map_messages( ) ) - ctp = ContentToProcess( - content_entries=[meta], - activity_metadata=activity_meta, - device_metadata=device_meta, - integrated_app_metadata=integrated_app, - protected_app_metadata=protected_app, - ) - req = ProcessContentRequest( - content_to_process=ctp, - user_id=resolved_user_id, # Use the resolved user_id for all messages - tenant_id=tenant_id, - correlation_id=meta.correlation_id, - process_inline=None, # Will be set based on execution mode - ) - results.append(req) + for index, purview_content in enumerate(mapped_contents): + content_entry = ProcessConversationMetadata( + identifier=message_id if index == 0 else f"{message_id}-{index}", + content=purview_content, + name=f"Agent Framework Message {message_id}", + is_truncated=False, + correlation_id=correlation_id, + sequence_number=base_sequence_number + index, + ) + ctp = ContentToProcess( + content_entry=content_entry, + activity_metadata=activity_meta, + device_metadata=device_meta, + integrated_app_metadata=integrated_app, + protected_app_metadata=protected_app, + ) + req = ProcessContentRequest( + content_to_process=ctp, + user_id=resolved_user_id, # Use the resolved user_id for all messages + tenant_id=tenant_id, + correlation_id=correlation_id, + process_inline=None, # Will be set based on execution mode + ) + results.append(req) return results, resolved_user_id async def _process_with_scopes(self, pc_request: ProcessContentRequest) -> ProcessContentResponse: @@ -231,7 +365,9 @@ async def _process_with_scopes(self, pc_request: ProcessContentRequest) -> Proce return await self._process_with_cached_scopes(pc_request, cached_ps_resp, cache_key) pc_request.process_inline = True - task = asyncio.create_task(self._refresh_protection_scopes_background(ps_req, cache_key, pc_request)) + # The background refresh gets its own copy: the foreground call mutates + # process_inline and scope_identifier on this request while that task runs. + task = asyncio.create_task(self._refresh_protection_scopes_background(ps_req, cache_key, copy(pc_request))) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) return await self._call_process_content(pc_request, cache_key, dlp_actions=[]) @@ -248,15 +384,23 @@ async def _process_with_cached_scopes( should_process, dlp_actions, execution_mode = self._check_applicable_scopes(pc_request, ps_resp) if should_process: - # Set process_inline based on execution mode - pc_request.process_inline = execution_mode == ExecutionMode.EVALUATE_INLINE + # Only an explicitly offline scope may skip inline evaluation. executionMode is an + # evolvable enum, so any unrecognised value must fail closed and be evaluated inline. + evaluate_offline = execution_mode == ExecutionMode.EVALUATE_OFFLINE + pc_request.process_inline = not evaluate_offline # If execution mode is offline, queue the PC request in background - if execution_mode != ExecutionMode.EVALUATE_INLINE: + if evaluate_offline: task = asyncio.create_task(self._process_content_background(pc_request, cache_key)) self._background_tasks.add(task) task.add_done_callback(self._background_tasks.discard) - return ProcessContentResponse(id="204", correlation_id=pc_request.correlation_id) + # Offline evaluation is asynchronous, but an explicit block action already known + # from the cached scopes must still be enforced rather than discarded. + return ProcessContentResponse( + id="204", + correlation_id=pc_request.correlation_id, + policy_actions=_blocking_actions(dlp_actions) or None, + ) return await self._call_process_content(pc_request, cache_key, dlp_actions=dlp_actions) @@ -373,15 +517,19 @@ def _check_applicable_scopes( loc.data_type and location.data_type and loc.data_type.lower().endswith(location.data_type.split(".")[-1].lower()) - and loc.value == location.value + and isinstance(loc.value, str) + and isinstance(location.value, str) + and _normalize_location_value(location.data_type, loc.value) + == _normalize_location_value(location.data_type, location.value) ): location_match = True break if activity_match and location_match: should_process = True - # If any scope has EvaluateInline, upgrade to inline mode - if scope.execution_mode == ExecutionMode.EVALUATE_INLINE: + # Only an explicitly offline scope may skip inline evaluation. execution_mode is an + # evolvable enum, so any unrecognised value is upgraded to inline (fail closed). + if scope.execution_mode != ExecutionMode.EVALUATE_OFFLINE: execution_mode = ExecutionMode.EVALUATE_INLINE if scope.policy_actions: diff --git a/python/packages/purview/tests/purview/conftest.py b/python/packages/purview/tests/purview/conftest.py index dc9a7024e8c..0089053b7c9 100644 --- a/python/packages/purview/tests/purview/conftest.py +++ b/python/packages/purview/tests/purview/conftest.py @@ -41,7 +41,7 @@ def _create_content(text: str = "Test") -> ContentToProcess: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) return ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, diff --git a/python/packages/purview/tests/purview/test_processor.py b/python/packages/purview/tests/purview/test_processor.py index 9d7388a416b..f5d2fe81efb 100644 --- a/python/packages/purview/tests/purview/test_processor.py +++ b/python/packages/purview/tests/purview/test_processor.py @@ -20,6 +20,8 @@ ProcessContentResponse, ProtectionScopeActivities, ProtectionScopeState, + PurviewBinaryContent, + PurviewTextContent, RestrictionAction, ) from agent_framework_purview._processor import ScopedContentProcessor, _is_valid_guid @@ -144,6 +146,126 @@ async def test_map_messages_creates_requests( assert requests[0].tenant_id == "12345678-1234-1234-1234-123456789012" assert user_id == "12345678-1234-1234-1234-123456789012" + async def test_map_messages_submits_each_content_item_in_separate_request( + self, processor: ScopedContentProcessor + ) -> None: + """Test _map_messages submits every content item in a separate request. + + Non-text content flattened to Message.text is empty for binary, tool-call and + tool-result content, which would have Purview classify an empty string. + """ + from agent_framework import Content + + secret = b"credit card 4532667785213500" + messages = [ + Message( + role="user", + contents=[ + Content.from_data(data=secret, media_type="application/octet-stream"), + Content( + "function_call", + call_id="call-1", + name="send_email", + arguments={"body": "ssn 120-98-1437"}, + ), + Content("function_result", call_id="call-1", result="account 999-12345"), + ], + ) + ] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert len(requests) == 3 + entries = [request.content_to_process.content_entries[0] for request in requests] + assert all(len(request.content_to_process.content_entries) == 1 for request in requests) + binary_entry, call_entry, result_entry = entries + assert isinstance(binary_entry.content, PurviewBinaryContent) + assert binary_entry.content.data == secret + + assert isinstance(call_entry.content, PurviewTextContent) + assert "send_email" in call_entry.content.data + assert "120-98-1437" in call_entry.content.data + + assert isinstance(result_entry.content, PurviewTextContent) + assert "999-12345" in result_entry.content.data + + # Every entry must carry real content; an empty payload would not be evaluated. + assert all(entry.content is not None for entry in entries) + assert not any(isinstance(entry.content, PurviewTextContent) and entry.content.data == "" for entry in entries) + + async def test_map_messages_skips_empty_text_and_data(self, processor: ScopedContentProcessor) -> None: + """Test _map_messages does not create requests for empty text or binary data.""" + from agent_framework import Content + + messages = [ + Message( + role="user", + contents=[ + "", + Content.from_data(data=b"", media_type="application/octet-stream"), + ], + ) + ] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert requests == [] + + async def test_map_messages_decodes_data_uri_with_media_type_parameters( + self, processor: ScopedContentProcessor + ) -> None: + """Test _map_messages decodes base64 data URIs whose media type carries parameters. + + A data URI may carry any number of ";parameter=value" segments between the media type and + ";base64". Failing to decode one leaves the payload as a base64 string that no classifier + can read. + """ + from agent_framework import Content + + secret = b"credit card 4532667785213500" + content = Content.from_data(data=secret, media_type="text/plain;charset=utf-8") + assert content.uri is not None + assert content.uri.startswith("data:text/plain;charset=utf-8;base64,") + + messages = [Message(role="user", contents=[content])] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + entries = requests[0].content_to_process.content_entries + assert len(entries) == 1 + assert isinstance(entries[0].content, PurviewBinaryContent) + assert entries[0].content.data == secret + + async def test_map_messages_submits_additional_properties_on_structured_content( + self, processor: ScopedContentProcessor + ) -> None: + """Test _map_messages serializes structured content whole, including additional_properties. + + Building an entry from name/arguments alone drops additional_properties, leaving a + data channel on tool calls and tool results that Purview never sees. + """ + from agent_framework import Content + + call = Content("function_call", call_id="call-1", name="send_email", arguments={"body": "hello"}) + call.additional_properties = {"hidden": "ssn 120-98-1437"} + result = Content("function_result", call_id="call-1", result="ok") + result.additional_properties = {"hidden": "card 4532667785213500"} + + messages = [Message(role="user", contents=[call, result])] + + requests, _ = await processor._map_messages(messages, Activity.UPLOAD_TEXT) + + assert len(requests) == 2 + call_entry = requests[0].content_to_process.content_entries[0] + result_entry = requests[1].content_to_process.content_entries[0] + + assert isinstance(call_entry.content, PurviewTextContent) + assert "120-98-1437" in call_entry.content.data + assert "send_email" in call_entry.content.data + + assert isinstance(result_entry.content, PurviewTextContent) + assert "4532667785213500" in result_entry.content.data + async def test_map_messages_without_defaults_gets_token_info(self, mock_client: AsyncMock) -> None: """Test _map_messages gets token info when settings lack some defaults.""" settings = PurviewSettings(app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012") @@ -220,6 +342,207 @@ async def test_check_applicable_scopes_with_block_action( assert len(actions) == 1 assert actions[0].action == DlpAction.BLOCK_ACCESS + async def test_check_applicable_scopes_matches_location_case_insensitively( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes matches location values regardless of GUID casing. + + A casing difference between the request location and the scope location would + otherwise hide an applicable block scope. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationApplication", + value="A1B2C3D4-E5F6-4A5B-8C9D-0E1F2A3B4C5D", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationApplication", + value="a1b2c3d4-e5f6-4a5b-8c9d-0e1f2a3b4c5d", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, execution_mode = ScopedContentProcessor._check_applicable_scopes( + pc_request, ps_response + ) + + assert should_process is True + assert execution_mode == ExecutionMode.EVALUATE_INLINE + assert dlp_actions + + async def test_check_applicable_scopes_matches_url_location_host_case_insensitively( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes ignores host casing on URL locations. + + The scheme and host of a URL are case-insensitive, so a casing difference there must not + hide an applicable scope. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="HTTPS://Contoso.com/sites/marketing", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com/sites/marketing", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is True + assert dlp_actions + + async def test_check_applicable_scopes_treats_url_location_path_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL path casing significant. + + URL paths are case-sensitive, so a scope scoped to one path must not match a different path + that differs only in casing. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://contoso.com/sites/Marketing", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com/sites/marketing", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + + async def test_check_applicable_scopes_treats_url_location_query_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL query casing significant. + + A query value is case-sensitive, and it may follow the host directly with no path between + them, so the authority has to end at the query delimiter as well as the path delimiter. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://contoso.com?label=Secret", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com?label=secret", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + + async def test_check_applicable_scopes_treats_url_location_fragment_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL fragment casing significant. + + A fragment may follow the host directly, so the authority has to end at the fragment + delimiter as well. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://contoso.com#Section", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://contoso.com#section", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + + async def test_check_applicable_scopes_treats_url_location_userinfo_as_case_sensitive( + self, process_content_request_factory + ) -> None: + """Test _check_applicable_scopes keeps URL userinfo casing significant. + + Credentials embedded in the authority are case-sensitive, so only the host half of the + authority may be folded. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + pc_request.content_to_process.protected_app_metadata.application_location = PolicyLocation( + data_type="microsoft.graph.policyLocationUrl", + value="https://alice:SecretPass@contoso.com/docs", + ) + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_INLINE, + locations=[ + PolicyLocation( + data_type="#microsoft.graph.policyLocationUrl", + value="https://alice:secretpass@contoso.com/docs", + ) + ], + policy_actions=[DlpActionInfo(action=DlpAction.BLOCK_ACCESS)], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + + should_process, dlp_actions, _ = ScopedContentProcessor._check_applicable_scopes(pc_request, ps_response) + + assert should_process is False + assert dlp_actions == [] + async def test_combine_policy_actions(self, processor: ScopedContentProcessor) -> None: """Test _combine_policy_actions merges action lists.""" action1 = DlpActionInfo(action=DlpAction.BLOCK_ACCESS, restrictionAction=RestrictionAction.BLOCK) @@ -290,6 +613,115 @@ async def test_process_with_scopes_calls_client_methods( mock_client.get_protection_scopes.assert_called_once() mock_client.send_content_activities.assert_called_once() + async def test_process_with_scopes_isolates_the_background_refresh_request( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test a cold scopes cache evaluates inline and hands the background refresh its own request. + + The background refresh must not observe foreground mutations of the request it was given. + """ + pc_request = process_content_request_factory() + cast(Any, processor._cache).get = AsyncMock(return_value=None) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + captured: list[Any] = [] + + async def capture_refresh(ps_req: Any, cache_key: str, request: Any) -> None: + captured.append(request) + + cast(Any, processor)._refresh_protection_scopes_background = capture_refresh + + await processor._process_with_scopes(pc_request) + await asyncio.gather(*list(processor._background_tasks)) + + # Content was evaluated inline, so there is no window in which it goes unevaluated. + mock_client.process_content.assert_called_once() + assert pc_request.process_inline is True + + # The background task got its own request object. + assert captured and captured[0] is not pc_request + + async def test_process_with_scopes_evaluates_unknown_execution_mode_inline( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test an unrecognised executionMode is evaluated inline. + + executionMode is an evolvable enum, so an unknown member must not skip enforcement. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.UNKNOWN_FUTURE_VALUE, + locations=[PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")], + policy_actions=[], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + cast(Any, processor._cache).get = AsyncMock(side_effect=[None, ps_response]) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + await processor._process_with_scopes(pc_request) + + mock_client.process_content.assert_called_once() + assert pc_request.process_inline is True + + async def test_process_with_scopes_reports_block_action_on_offline_scope( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test a block verdict known from an offline scope is still reported rather than discarded.""" + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_OFFLINE, + locations=[PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")], + policy_actions=[ + DlpActionInfo(action=DlpAction.RESTRICT_ACCESS, restriction_action=RestrictionAction.BLOCK) + ], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + cast(Any, processor._cache).get = AsyncMock(side_effect=[None, ps_response]) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + response = await processor._process_with_scopes(pc_request) + await asyncio.gather(*list(processor._background_tasks)) + + assert response.policy_actions + assert all( + action.action == DlpAction.RESTRICT_ACCESS and action.restriction_action == RestrictionAction.BLOCK + for action in response.policy_actions + ) + + async def test_process_with_scopes_ignores_non_blocking_restriction_on_offline_scope( + self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory + ) -> None: + """Test a restrictAccess scope whose restriction mode does not block is not reported as blocking. + + restrictAccess carries a separate restriction mode which may be audit, warn or allow. Only an + explicit block mode withholds the content, so the other modes must not surface as a verdict. + """ + from agent_framework_purview._models import ProtectionScopesResponse + + pc_request = process_content_request_factory() + scope = PolicyScope( + activities=ProtectionScopeActivities.UPLOAD_TEXT, + execution_mode=ExecutionMode.EVALUATE_OFFLINE, + locations=[PolicyLocation(data_type="microsoft.graph.policyLocationApplication", value="app-id")], + policy_actions=[ + DlpActionInfo(action=DlpAction.RESTRICT_ACCESS, restriction_action=RestrictionAction.OTHER) + ], + ) + ps_response = ProtectionScopesResponse(scopes=[scope]) + cast(Any, processor._cache).get = AsyncMock(side_effect=[None, ps_response]) + mock_client.process_content.return_value = ProcessContentResponse(id="1") + + response = await processor._process_with_scopes(pc_request) + await asyncio.gather(*list(processor._background_tasks)) + + assert not response.policy_actions + async def test_process_with_scopes_preserves_restriction_only_policy_actions( self, processor: ScopedContentProcessor, mock_client: AsyncMock, process_content_request_factory ) -> None: @@ -516,8 +948,8 @@ async def test_map_messages_with_provided_user_id_fallback(self, mock_client: As assert user_id == "32345678-1234-1234-1234-123456789012" assert requests[0].user_id == "32345678-1234-1234-1234-123456789012" - async def test_map_messages_returns_empty_when_no_user_id(self, mock_client: AsyncMock) -> None: - """Test that empty results are returned when user_id cannot be resolved.""" + async def test_map_messages_raises_when_no_user_id(self, mock_client: AsyncMock) -> None: + """Test that an unresolvable user_id fails closed instead of skipping evaluation.""" settings = PurviewSettings( app_name="Test App", tenant_id="12345678-1234-1234-1234-123456789012", @@ -533,10 +965,8 @@ async def test_map_messages_returns_empty_when_no_user_id(self, mock_client: Asy messages = [Message(role="user", contents=["Test message"])] - requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) - - assert len(requests) == 0 - assert user_id is None + with pytest.raises(ValueError, match="No user id"): + await processor._map_messages(messages, Activity.UPLOAD_TEXT) async def test_process_content_sends_activities_when_not_applicable( self, mock_client: AsyncMock, process_content_request_factory @@ -705,11 +1135,9 @@ async def test_author_name_ignored_if_not_valid_guid( ) ] - requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT) - - # Should return empty since author_name is not a valid GUID - assert user_id is None - assert len(requests) == 0 + # author_name is not a valid GUID, so no identity resolves and evaluation fails closed + with pytest.raises(ValueError, match="No user id"): + await processor._map_messages(messages, Activity.UPLOAD_TEXT) async def test_provided_user_id_used_as_last_resort( self, mock_client: AsyncMock, settings: PurviewSettings @@ -739,10 +1167,29 @@ async def test_invalid_provided_user_id_ignored(self, mock_client: AsyncMock, se messages = [Message(role="user", contents=["Test"])] - requests, user_id = await processor._map_messages(messages, Activity.UPLOAD_TEXT, provided_user_id="not-a-guid") + with pytest.raises(ValueError, match="No user id"): + await processor._map_messages(messages, Activity.UPLOAD_TEXT, provided_user_id="not-a-guid") - assert user_id is None - assert len(requests) == 0 + async def test_unresolvable_user_id_blocks_processing( + self, mock_client: AsyncMock, settings: PurviewSettings + ) -> None: + """Test process_messages raises rather than sending content when no user id resolves. + + Without a user id there is no policy to evaluate against, so the content must not be + forwarded to Purview or on to the model. + """ + mock_client.get_user_info_from_token.return_value = { + "tenant_id": "12345678-1234-1234-1234-123456789012", + "client_id": "12345678-1234-1234-1234-123456789012", + } + processor = ScopedContentProcessor(mock_client, settings) + + messages = [Message(role="user", contents=["ssn 120-98-1437"])] + + with pytest.raises(ValueError, match="No user id"): + await processor.process_messages(messages, Activity.UPLOAD_TEXT) + + mock_client.process_content.assert_not_called() async def test_multiple_messages_same_user_id(self, mock_client: AsyncMock, settings: PurviewSettings) -> None: """Test that all messages use the same resolved user_id.""" diff --git a/python/packages/purview/tests/purview/test_purview_models.py b/python/packages/purview/tests/purview/test_purview_models.py index 4af581014d9..3e3e3fb929a 100644 --- a/python/packages/purview/tests/purview/test_purview_models.py +++ b/python/packages/purview/tests/purview/test_purview_models.py @@ -91,7 +91,7 @@ def test_content_to_process_with_nested_structures(self) -> None: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) content = ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, @@ -188,7 +188,7 @@ def test_content_serialization_uses_aliases(self) -> None: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) content = ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, @@ -198,6 +198,7 @@ def test_content_serialization_uses_aliases(self) -> None: dumped = content.model_dump(by_alias=True, exclude_none=True, mode="json") assert "contentEntries" in dumped + assert len(dumped["contentEntries"]) == 1 assert "activityMetadata" in dumped assert "deviceMetadata" in dumped assert "integratedAppMetadata" in dumped @@ -224,7 +225,7 @@ def test_process_content_request_excludes_private_fields(self) -> None: protected_app = ProtectedAppMetadata(name="Protected", version="1.0", application_location=location) content = ContentToProcess( - content_entries=[metadata], + content_entry=metadata, activity_metadata=activity_meta, device_metadata=device_meta, integrated_app_metadata=integrated_app, diff --git a/python/samples/02-agents/context_providers/README.md b/python/samples/02-agents/context_providers/README.md index df0a3a3b1cc..1c8b5990146 100644 --- a/python/samples/02-agents/context_providers/README.md +++ b/python/samples/02-agents/context_providers/README.md @@ -57,3 +57,30 @@ These samples demonstrate how to use context providers to enrich agent conversat - Azure CLI authentication (`az login`) See each subfolder's README for provider-specific prerequisites. + +## Application-controlled modes + +Configure `AgentModeProvider(expose_mode_set=False)` to hide the built-in setter +when your application owns mode changes. Use `expose_mode_get=False` to hide the +getter independently, or set both flags to `False` to expose neither tool. Both +flags default to `True`. Mode state, per-turn workflow instructions, and external +mode-change notifications remain active even when both tools are hidden. + +Pass the configured provider through `create_harness_agent(mode_provider=...)` +(or `Agent(context_providers=[...])`). Supply any replacement tool, such as +`update_mode`, through the existing `tools` argument. Keep that tool and your UI +on the same session-backed state by using `get_agent_mode` and `set_agent_mode` +with the provider's `source_id` and `available_modes`, and its `default_mode` when +reading. A replacement tool should call `set_agent_mode(..., notify=False)` +because the agent already observes the tool result; this also clears any pending +external-change notification. UI-driven changes retain the default `notify=True` +so the agent sees the external change on its next run. Do not disable the entire +provider with `disable_mode=True`. + +Built-in guidance only advertises enabled tools. Without the built-in setter, +it defers approved transitions to the application's configured mode-change +mechanism. Supply replacement-specific guidance through the existing +`instructions` and `mode_instructions` options; caller-supplied text is not +rewritten when a tool is hidden, and existing placeholders still expand. +These flags only omit tools contributed by this provider: they do not filter +application-supplied tools or prevent application code from changing mode state. diff --git a/python/samples/02-agents/observability/README.md b/python/samples/02-agents/observability/README.md index 822006b17dd..e69af64f684 100644 --- a/python/samples/02-agents/observability/README.md +++ b/python/samples/02-agents/observability/README.md @@ -88,6 +88,18 @@ from microsoft.opentelemetry import use_microsoft_opentelemetry use_microsoft_opentelemetry(enable_azure_monitor=True) ``` +To disable Agent Framework's baseline GenAI message events without changing providers or exporters configured by a third party, use the instrumentation-only API: + +```python +from agent_framework.observability import enable_instrumentation + +enable_instrumentation(enable_message_events=False) +``` + +An explicit `True` or `False` overrides the current message-event setting, including a value read from `ENABLE_MESSAGE_EVENTS`. Omitting the argument or passing `None` preserves the current setting without re-reading the environment, including settings established by `configure_otel_providers()` or an earlier explicit call. + +This flag controls baseline v1.36.0 events such as `gen_ai.user.message` and `gen_ai.choice`; it does not disable experimental message span attributes, spans, or metrics. Message events still require sensitive-data capture to be enabled separately. Existing `enable_sensitive_data` behavior is unchanged, and [sticky disable](#disabling-instrumentation) still requires `force=True` to re-enable instrumentation. + ```python from azure.monitor.opentelemetry import configure_azure_monitor from agent_framework.observability import create_resource, enable_sensitive_telemetry diff --git a/python/uv.lock b/python/uv.lock index 27bf49e5eb1..2ae5de419bd 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -449,6 +449,7 @@ dependencies = [ { name = "opentelemetry-api", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "pydantic", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "python-dotenv", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] @@ -491,6 +492,7 @@ dev = [ { name = "azure-ai-agentserver-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "graphviz", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "opentelemetry-exporter-otlp", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, + { name = "types-pyyaml", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] [package.metadata] @@ -529,6 +531,7 @@ requires-dist = [ { name = "opentelemetry-api", specifier = ">=1.39.0,<2" }, { name = "pydantic", specifier = ">=2,<3" }, { name = "python-dotenv", specifier = ">=1,<2" }, + { name = "pyyaml", specifier = ">=6.0,<7.0" }, { name = "typing-extensions", specifier = ">=4.15.0,<5" }, ] provides-extras = ["all"] @@ -538,6 +541,7 @@ dev = [ { name = "azure-ai-agentserver-core", specifier = ">=2.0.0b7,<3" }, { name = "graphviz", specifier = ">=0.20,<1" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.39.0,<2" }, + { name = "types-pyyaml", specifier = ">=6.0,<7.0" }, ] [[package]]