Fix MCP serialization for telemetry. Add unit tests. (#42739)

* Fix MCP serialization for telemetry. Add unit tests.

* Fix

Author: nick863

sdk/ai/azure-ai-agents/CHANGELOG.md

@@ -12,6 +12,7 @@
 
 * Fix the issue with logging Agent message, when the message has "in progress" status (related to [issue](https://github.com/Azure/azure-sdk-for-python/issues/42645)).
 * Fix the issue with `RunStepOpenAPIToolCall` logging [issue](https://github.com/Azure/azure-sdk-for-python/issues/42645).
+* Fix the issue with `RunStepMcpToolCall` logging [issue](https://github.com/Azure/azure-sdk-for-python/issues/42689).
 
 ### Sample updates
 

sdk/ai/azure-ai-agents/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ai/azure-ai-agents",
-  "Tag": "python/ai/azure-ai-agents_ac2e010ef5"
+  "Tag": "python/ai/azure-ai-agents_28d79451bb"
 }

sdk/ai/azure-ai-agents/azure/ai/agents/telemetry/_ai_agents_instrumentor.py

@@ -30,6 +30,7 @@
     RunStepDeltaChunk,
     RunStepError,
     RunStepFunctionToolCall,
+    RunStepMcpToolCall,
     RunStepOpenAPIToolCall,
     RunStepToolCallDetails,
     RunStepCodeInterpreterToolCall,
@@ -448,6 +449,15 @@ def _process_tool_calls(self, step: RunStep) -> List[Dict[str, Any]]:
                     "type": t.type,
                     'function': t.as_dict().get('function', {})
                 }
+            elif isinstance(t, RunStepMcpToolCall):
+                tool_call = {
+                    "id": t.id,
+                    "type": t.type,
+                    "arguments": t.arguments,
+                    "name": t.name,
+                    "output": t.output,
+                    "server_label": t.server_label or ""
+                }
             else:
                 tool_details = t.as_dict()[t.type]
 
@@ -2066,7 +2076,10 @@ def on_thread_message(self, message: "ThreadMessage") -> None:  # type: ignore[f
         else:
             retval = super().on_thread_message(message)  # pylint: disable=assignment-from-none # type: ignore
 
-        # Message status may be in progress, even if the thread.message.completed event has arrived.
+        # TODO: Workaround for issue where message.status may be IN_PROGRESS even after a thread.message.completed event has arrived.
+        # See work item 4636616 and 4636299 for details.
+        # When the work item is resolved, change this code back to:
+        # if message.status in {MessageStatus.COMPLETED, MessageStatus.INCOMPLETE}
         if message.status in {MessageStatus.COMPLETED, MessageStatus.INCOMPLETE} or (message.status == MessageStatus.IN_PROGRESS and message.content):
             self.last_message = message
 
@@ -2204,7 +2217,10 @@ async def on_thread_message(self, message: "ThreadMessage") -> None:  # type: ig
         else:
             retval = await super().on_thread_message(message)  # type: ignore
 
-        # Message status may be in progress, even if the thread.message.completed event has arrived.
+        # TODO: Workaround for issue where message.status may be IN_PROGRESS even after a thread.message.completed event has arrived.
+        # See work item 4636616 and 4636299 for details.
+        # When the work item is resolved, change this code back to:
+        # if message.status in {MessageStatus.COMPLETED, MessageStatus.INCOMPLETE}
         if message.status in {MessageStatus.COMPLETED, MessageStatus.INCOMPLETE} or (message.status == MessageStatus.IN_PROGRESS and message.content):
             self.last_message = message
 

sdk/ai/azure-ai-agents/tests/gen_ai_trace_verifier.py

@@ -21,51 +21,44 @@ def check_span_attributes(self, span, attributes):
         for attribute_name in span.attributes.keys():
             # Check if the attribute name exists in the input attributes
             if attribute_name not in attribute_dict:
-                print("Attribute name " + attribute_name + " not in attribute dictionary")
-                return False
+                raise AssertionError("Attribute name " + attribute_name + " not in attribute dictionary")
 
             attribute_value = attribute_dict[attribute_name]
             if isinstance(attribute_value, list):
                 # Check if the attribute value in the span matches the provided list
                 if span.attributes[attribute_name] != attribute_value:
-                    print(
+                    raise AssertionError(
                         "Attribute value list "
                         + str(span.attributes[attribute_name])
                         + " does not match with "
                         + str(attribute_value)
                     )
-                    return False
             elif isinstance(attribute_value, tuple):
                 # Check if the attribute value in the span matches the provided list
                 if span.attributes[attribute_name] != attribute_value:
-                    print(
+                    raise AssertionError(
                         "Attribute value tuple "
                         + str(span.attributes[attribute_name])
                         + " does not match with "
                         + str(attribute_value)
                     )
-                    return False
             else:
                 # Check if the attribute value matches the provided value
                 if attribute_value == "+":
                     if not isinstance(span.attributes[attribute_name], numbers.Number):
-                        print("Attribute value " + str(span.attributes[attribute_name]) + " is not a number")
-                        return False
+                        raise AssertionError("Attribute value " + str(span.attributes[attribute_name]) + " is not a number")
                     if span.attributes[attribute_name] < 0:
-                        print("Attribute value " + str(span.attributes[attribute_name]) + " is negative")
-                        return False
+                        raise AssertionError("Attribute value " + str(span.attributes[attribute_name]) + " is negative")
                 elif attribute_value != "" and span.attributes[attribute_name] != attribute_value:
-                    print(
+                    raise AssertionError(
                         "Attribute value "
                         + str(span.attributes[attribute_name])
                         + " does not match with "
                         + str(attribute_value)
                     )
-                    return False
                 # Check if the attribute value in the span is not empty when the provided value is ""
                 elif attribute_value == "" and not span.attributes[attribute_name]:
-                    print("Expected non-empty attribute value")
-                    return False
+                    raise AssertionError("Expected non-empty attribute value")
 
         return True
 
@@ -76,63 +69,53 @@ def check_decorator_span_attributes(self, span: Span, attributes: List[tuple]) -
         # Ensure all required attributes are present in the span
         for attribute_name in attribute_dict.keys():
             if attribute_name not in span.attributes:
-                print("Required attribute name " + attribute_name + " not found in span attributes")
-                return False
+                raise AssertionError("Required attribute name " + attribute_name + " not found in span attributes")
 
         for attribute_name in span.attributes.keys():
             # Check if the attribute name exists in the input attributes
             if attribute_name not in attribute_dict:
-                print("Attribute name " + attribute_name + " not in attribute dictionary")
-                return False
+                raise AssertionError("Attribute name " + attribute_name + " not in attribute dictionary")
 
             attribute_value = attribute_dict[attribute_name]
             span_value = span.attributes[attribute_name]
 
             if isinstance(attribute_value, (list, tuple)):
                 # Convert both to lists for comparison
                 if list(span_value) != list(attribute_value):
-                    print(
+                    raise AssertionError(
                         "Attribute value list/tuple " + str(span_value) + " does not match with " + str(attribute_value)
                     )
-                    return False
             elif isinstance(attribute_value, dict):
                 # Check if both are dictionaries and compare them
                 if not isinstance(span_value, dict) or span_value != attribute_value:
-                    print("Attribute value dict " + str(span_value) + " does not match with " + str(attribute_value))
-                    return False
+                    raise AssertionError("Attribute value dict " + str(span_value) + " does not match with " + str(attribute_value))
             else:
                 # Check if the attribute value matches the provided value
                 if attribute_value == "+":
                     if not isinstance(span_value, numbers.Number):
-                        print("Attribute value " + str(span_value) + " is not a number")
-                        return False
+                        raise AssertionError("Attribute value " + str(span_value) + " is not a number")
                     if span_value < 0:
-                        print("Attribute value " + str(span_value) + " is negative")
-                        return False
+                        raise AssertionError("Attribute value " + str(span_value) + " is negative")
                 elif attribute_value != "" and span_value != attribute_value:
-                    print("Attribute value " + str(span_value) + " does not match with " + str(attribute_value))
-                    return False
+                    raise AssertionError("Attribute value " + str(span_value) + " does not match with " + str(attribute_value))
                 # Check if the attribute value in the span is not empty when the provided value is ""
                 elif attribute_value == "" and not span_value:
-                    print("Expected non-empty attribute value")
-                    return False
+                    raise AssertionError("Expected non-empty attribute value")
 
         return True
 
     def is_valid_json(self, my_string):
         try:
-            json_object = json.loads(my_string)
+            json.loads(my_string)
         except ValueError as e1:
             return False
         except TypeError as e2:
             return False
         return True
 
     def check_json_string(self, expected_json, actual_json):
-        if self.is_valid_json(expected_json) and self.is_valid_json(actual_json):
-            return self.check_event_attributes(json.loads(expected_json), json.loads(actual_json))
-        else:
-            return False
+        assert self.is_valid_json(expected_json) and self.is_valid_json(actual_json)
+        return self.check_event_attributes(json.loads(expected_json), json.loads(actual_json))
 
     def check_event_attributes(self, expected_dict, actual_dict):
         if set(expected_dict.keys()) != set(actual_dict.keys()):
@@ -144,56 +127,39 @@ def check_event_attributes(self, expected_dict, actual_dict):
                 actual_val = json.dumps(actual_dict)
             else:
                 actual_val = actual_dict
-            print("check_event_attributes: keys do not match: " + expected_val + "!=" + actual_val)
-            return False
+            raise AssertionError("check_event_attributes: keys do not match: " + expected_val + "!=" + actual_val)
         for key, expected_val in expected_dict.items():
             if key not in actual_dict:
-                print("check_event_attributes: key not found")
-                return False
+                raise AssertionError(f"check_event_attributes: key {key} not found in actuals")
             actual_val = actual_dict[key]
 
             if self.is_valid_json(expected_val):
                 if not self.is_valid_json(actual_val):
-                    print("check_event_attributes: actual_val is not valid json")
-                    return False
-                if not self.check_json_string(expected_val, actual_val):
-                    print("check_event_attributes: check_json_string failed")
-                    return False
+                    raise AssertionError(f"check_event_attributes: actual_val for {key} is not valid json")
+                self.check_json_string(expected_val, actual_val)
             elif isinstance(expected_val, dict):
                 if not isinstance(actual_val, dict):
-                    print("check_event_attributes: actual_val is not dict")
-                    return False
-                if not self.check_event_attributes(expected_val, actual_val):
-                    print("check_event_attributes: check_event_attributes failed")
-                    return False
+                    raise AssertionError(f"check_event_attributes: actual_val for {key} is not dict")
+                self.check_event_attributes(expected_val, actual_val)
             elif isinstance(expected_val, list):
                 if not isinstance(actual_val, list):
-                    print("check_event_attributes: actual_val is not list")
-                    return False
+                    raise AssertionError(f"check_event_attributes: actual_val for {key}  is not list")
                 if len(expected_val) != len(actual_val):
-                    print("check_event_attributes: list lengths do not match")
-                    return False
+                    raise AssertionError(f"check_event_attributes: list lengths do not match for key {key}: expected {len(expected_val)}, actual {len(actual_val)}")
                 for expected_list, actual_list in zip(expected_val, actual_val):
-                    if not self.check_event_attributes(expected_list, actual_list):
-                        print("check_event_attributes: check_event_attributes for list failed")
-                        return False
+                    self.check_event_attributes(expected_list, actual_list)
             elif isinstance(expected_val, str) and expected_val == "*":
                 if actual_val == "":
-                    print("check_event_attributes: actual_val is empty")
-                    return False
+                    raise AssertionError(f"check_event_attributes: actual_val for {key} is empty")
             elif isinstance(expected_val, str) and expected_val == "+":
-                if not isinstance(actual_val, numbers.Number):
-                    return False
-                if actual_val < 0:
-                    return False
+                assert isinstance(actual_val, numbers.Number), f"The {key} is not a number."
+                assert actual_val > 0, f"The {key} is <0 {actual_val}" 
             elif expected_val != actual_val:
                 if isinstance(expected_val, dict):
                     expected_val = json.dumps(expected_val)
                 if isinstance(actual_val, dict):
                     actual_val = json.dumps(actual_val)
-                print("check_event_attributes: values do not match: " + expected_val + "!=" + actual_val)
-                return False
-        return True
+                raise AssertionError(f"check_event_attributes: values do not match for key {key}: {expected_val} != {actual_val}")
 
     def check_span_events(self, span, expected_events):
         print("Checking span: " + span.name)
@@ -202,17 +168,13 @@ def check_span_events(self, span, expected_events):
         for expected_event in expected_events:
             for actual_event in span_events:
                 if expected_event["name"] == actual_event.name:
-                    if not self.check_event_attributes(expected_event["attributes"], actual_event.attributes._dict):
-                        print("Event attributes do not match")
-                        return False
+                    self.check_event_attributes(expected_event["attributes"], actual_event.attributes._dict)
                     span_events.remove(actual_event)  # Remove the matched event from the span_events
                     break
             else:
-                print("check_span_events: event not found")
-                return False  # If no match found for an expected event
+                raise AssertionError("check_span_events: event not found")
 
         if len(span_events) > 0:  # If there are any additional events in the span_events
-            print("check_span_events: unexpected event found")
-            return False
+            raise AssertionError("check_span_events: unexpected event found")
 
         return True