.NET: Fix Usage and Annotations not present in AG-UI events (#3752)

dotnet/src/Microsoft.Agents.AI.AGUI/Shared/AGUIEventTypes.cs

@@ -31,4 +31,6 @@ internal static class AGUIEventTypes
     public const string StateSnapshot = "STATE_SNAPSHOT";
 
     public const string StateDelta = "STATE_DELTA";
+
+    public const string Custom = "CUSTOM";
 }

dotnet/src/Microsoft.Agents.AI.AGUI/Shared/AGUIJsonSerializerContext.cs

@@ -46,6 +46,7 @@ namespace Microsoft.Agents.AI.AGUI;
 [JsonSerializable(typeof(ToolCallResultEvent))]
 [JsonSerializable(typeof(StateSnapshotEvent))]
 [JsonSerializable(typeof(StateDeltaEvent))]
+[JsonSerializable(typeof(CustomEvent))]
 [JsonSerializable(typeof(IDictionary<string, object?>))]
 [JsonSerializable(typeof(Dictionary<string, object?>))]
 [JsonSerializable(typeof(IDictionary<string, System.Text.Json.JsonElement?>))]

dotnet/src/Microsoft.Agents.AI.AGUI/Shared/BaseEventJsonConverter.cs

@@ -47,6 +47,7 @@ public override BaseEvent Read(
             AGUIEventTypes.ToolCallEnd => jsonElement.Deserialize(options.GetTypeInfo(typeof(ToolCallEndEvent))) as ToolCallEndEvent,
             AGUIEventTypes.ToolCallResult => jsonElement.Deserialize(options.GetTypeInfo(typeof(ToolCallResultEvent))) as ToolCallResultEvent,
             AGUIEventTypes.StateSnapshot => jsonElement.Deserialize(options.GetTypeInfo(typeof(StateSnapshotEvent))) as StateSnapshotEvent,
+            AGUIEventTypes.Custom => jsonElement.Deserialize(options.GetTypeInfo(typeof(CustomEvent))) as CustomEvent,
             _ => throw new JsonException($"Unknown BaseEvent type discriminator: '{discriminator}'")
         };
 
@@ -102,6 +103,9 @@ public override void Write(
             case StateDeltaEvent stateDelta:
                 JsonSerializer.Serialize(writer, stateDelta, options.GetTypeInfo(typeof(StateDeltaEvent)));
                 break;
+            case CustomEvent customEvent:
+                JsonSerializer.Serialize(writer, customEvent, options.GetTypeInfo(typeof(CustomEvent)));
+                break;
             default:
                 throw new InvalidOperationException($"Unknown event type: {value.GetType().Name}");
         }

dotnet/src/Microsoft.Agents.AI.AGUI/Shared/ChatResponseUpdateAGUIExtensions.cs

@@ -375,6 +375,12 @@ chatResponse.Contents[0] is TextContent &&
                     MessageId = chatResponse.MessageId!,
                     Delta = textContent.Text
                 };
+
+                // Emit annotations if present on the text content
+                if (textContent.Annotations is { Count: > 0 })
+                {
+                    yield return CreateAnnotationsCustomEvent(textContent.Annotations, jsonSerializerOptions);
+                }
             }
 
             // Emit tool call events and tool result events
@@ -463,6 +469,11 @@ chatResponse.Contents[0] is TextContent &&
                             };
                         }
                     }
+                    else if (content is UsageContent usageContent)
+                    {
+                        // Emit usage data as a custom event
+                        yield return CreateUsageCustomEvent(usageContent, jsonSerializerOptions);
+                    }
                 }
             }
         }
@@ -493,4 +504,77 @@ chatResponse.Contents[0] is TextContent &&
             _ => JsonSerializer.Serialize(functionResultContent.Result, options.GetTypeInfo(functionResultContent.Result.GetType())),
         };
     }
+
+    private static CustomEvent CreateUsageCustomEvent(UsageContent usageContent, JsonSerializerOptions jsonSerializerOptions)
+    {
+        using var buffer = new System.IO.MemoryStream();
+        using (var writer = new Utf8JsonWriter(buffer))
+        {
+            writer.WriteStartObject();
+            writer.WriteNumber("inputTokenCount", usageContent.Details.InputTokenCount ?? 0);
+            writer.WriteNumber("outputTokenCount", usageContent.Details.OutputTokenCount ?? 0);
+            writer.WriteNumber("totalTokenCount", usageContent.Details.TotalTokenCount ?? 0);
+            writer.WriteEndObject();
+        }
+
+        return new CustomEvent
+        {
+            Name = "usage",
+            Value = JsonSerializer.Deserialize(buffer.ToArray(), jsonSerializerOptions.GetTypeInfo(typeof(JsonElement))) as JsonElement?
+        };
+    }
+
+    private static CustomEvent CreateAnnotationsCustomEvent(
+        IList<AIAnnotation> annotations,
+        JsonSerializerOptions jsonSerializerOptions)
+    {
+        using var buffer = new System.IO.MemoryStream();
+        using (var writer = new Utf8JsonWriter(buffer))
+        {
+            writer.WriteStartArray();
+            foreach (var annotation in annotations)
+            {
+                writer.WriteStartObject();
+                writer.WriteString("type", annotation.GetType().Name);
+
+                if (annotation is CitationAnnotation citation)
+                {
+                    if (citation.Title is not null)
+                    {
+                        writer.WriteString("title", citation.Title);
+                    }
+
+                    if (citation.Url is not null)
+                    {
+                        writer.WriteString("url", citation.Url.ToString());
+                    }
+
+                    if (citation.FileId is not null)
+                    {
+                        writer.WriteString("fileId", citation.FileId);
+                    }
+
+                    if (citation.ToolName is not null)
+                    {
+                        writer.WriteString("toolName", citation.ToolName);
+                    }
+
+                    if (citation.Snippet is not null)
+                    {
+                        writer.WriteString("snippet", citation.Snippet);
+                    }
+                }
+
+                writer.WriteEndObject();
+            }
+
+            writer.WriteEndArray();
+        }
+
+        return new CustomEvent
+        {
+            Name = "annotations",
+            Value = JsonSerializer.Deserialize(buffer.ToArray(), jsonSerializerOptions.GetTypeInfo(typeof(JsonElement))) as JsonElement?
+        };
+    }
 }

dotnet/src/Microsoft.Agents.AI.AGUI/Shared/CustomEvent.cs

@@ -0,0 +1,24 @@
+// Copyright (c) Microsoft. All rights reserved.
+
+using System.Text.Json;
+using System.Text.Json.Serialization;
+
+#if ASPNETCORE
+namespace Microsoft.Agents.AI.Hosting.AGUI.AspNetCore.Shared;
+#else
+namespace Microsoft.Agents.AI.AGUI.Shared;
+#endif
+
+internal sealed class CustomEvent : BaseEvent
+{
+    public CustomEvent()
+    {
+        this.Type = AGUIEventTypes.Custom;
+    }
+
+    [JsonPropertyName("name")]
+    public string Name { get; set; } = string.Empty;
+
+    [JsonPropertyName("value")]
+    public JsonElement? Value { get; set; }
+}