feat: Support optional spanKind for InvokeAgentScope and ExecuteToolScope (#195)

* Initial plan

* feat: add optional span_kind parameter to InvokeAgentScope and ExecuteToolScope

Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com>

* fix test and lint errors

---------

Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Co-authored-by: nikhilNava <211831449+nikhilNava@users.noreply.github.com>
Co-authored-by: Nikhil Navakiran <nikhil.navakiran@gmail.com>
Co-authored-by: Nikhil Chitlur Navakiran (from Dev Box) <nikhilc@microsoft.com>

Author: 3 people

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/execute_tool_scope.py

@@ -3,6 +3,8 @@
 
 from datetime import datetime
 
+from opentelemetry.trace import SpanKind
+
 from .agent_details import AgentDetails
 from .constants import (
     CHANNEL_LINK_KEY,
@@ -34,6 +36,7 @@ def start(
         parent_id: str | None = None,
         start_time: datetime | None = None,
         end_time: datetime | None = None,
+        span_kind: SpanKind | None = None,
     ) -> "ExecuteToolScope":
         """Creates and starts a new scope for tool execution tracing.
 
@@ -49,12 +52,21 @@ def start(
             end_time: Optional explicit end time as a datetime object. When provided,
                 the span will use this timestamp when disposed instead of the
                 current wall-clock time.
+            span_kind: Optional span kind override. Defaults to ``SpanKind.INTERNAL``.
+                Use ``SpanKind.CLIENT`` when the tool calls an external service.
 
         Returns:
             A new ExecuteToolScope instance
         """
         return ExecuteToolScope(
-            details, agent_details, tenant_details, request, parent_id, start_time, end_time
+            details,
+            agent_details,
+            tenant_details,
+            request,
+            parent_id,
+            start_time,
+            end_time,
+            span_kind,
         )
 
     def __init__(
@@ -66,6 +78,7 @@ def __init__(
         parent_id: str | None = None,
         start_time: datetime | None = None,
         end_time: datetime | None = None,
+        span_kind: SpanKind | None = None,
     ):
         """Initialize the tool execution scope.
 
@@ -81,9 +94,11 @@ def __init__(
             end_time: Optional explicit end time as a datetime object. When provided,
                 the span will use this timestamp when disposed instead of the
                 current wall-clock time.
+            span_kind: Optional span kind override. Defaults to ``SpanKind.INTERNAL``.
+                Use ``SpanKind.CLIENT`` when the tool calls an external service.
         """
         super().__init__(
-            kind="Internal",
+            kind=span_kind if span_kind is not None else SpanKind.INTERNAL,
             operation_name=EXECUTE_TOOL_OPERATION_NAME,
             activity_name=f"{EXECUTE_TOOL_OPERATION_NAME} {details.tool_name}",
             agent_details=agent_details,

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/invoke_agent_scope.py

@@ -6,6 +6,8 @@
 import logging
 from datetime import datetime
 
+from opentelemetry.trace import SpanKind
+
 from .agent_details import AgentDetails
 from .constants import (
     CHANNEL_LINK_KEY,
@@ -50,6 +52,7 @@ def start(
         caller_details: CallerDetails | None = None,
         start_time: datetime | None = None,
         end_time: datetime | None = None,
+        span_kind: SpanKind | None = None,
     ) -> "InvokeAgentScope":
         """Create and start a new scope for agent invocation tracing.
 
@@ -62,6 +65,8 @@ def start(
             caller_details: Optional details of the non-agentic caller
             start_time: Optional explicit start time as a datetime object.
             end_time: Optional explicit end time as a datetime object.
+            span_kind: Optional span kind override. Defaults to ``SpanKind.CLIENT``.
+                Use ``SpanKind.SERVER`` when the agent is receiving an inbound request.
 
         Returns:
             A new InvokeAgentScope instance
@@ -74,6 +79,7 @@ def start(
             caller_details,
             start_time,
             end_time,
+            span_kind,
         )
 
     def __init__(
@@ -85,6 +91,7 @@ def __init__(
         caller_details: CallerDetails | None = None,
         start_time: datetime | None = None,
         end_time: datetime | None = None,
+        span_kind: SpanKind | None = None,
     ):
         """Initialize the agent invocation scope.
 
@@ -96,6 +103,8 @@ def __init__(
             caller_details: Optional details of the non-agentic caller
             start_time: Optional explicit start time as a datetime object.
             end_time: Optional explicit end time as a datetime object.
+            span_kind: Optional span kind override. Defaults to ``SpanKind.CLIENT``.
+                Use ``SpanKind.SERVER`` when the agent is receiving an inbound request.
         """
         activity_name = INVOKE_AGENT_OPERATION_NAME
         if invoke_agent_details.details.agent_name:
@@ -104,7 +113,7 @@ def __init__(
             )
 
         super().__init__(
-            kind="Client",
+            kind=span_kind if span_kind is not None else SpanKind.CLIENT,
             operation_name=INVOKE_AGENT_OPERATION_NAME,
             activity_name=activity_name,
             agent_details=invoke_agent_details.details,

libraries/microsoft-agents-a365-observability-core/microsoft_agents_a365/observability/core/opentelemetry_scope.py

@@ -92,7 +92,7 @@ def _datetime_to_ns(dt: datetime | None) -> int | None:
 
     def __init__(
         self,
-        kind: str,
+        kind: "str | SpanKind",
         operation_name: str,
         activity_name: str,
         agent_details: "AgentDetails | None" = None,
@@ -104,7 +104,9 @@ def __init__(
         """Initialize the OpenTelemetry scope.
 
         Args:
-            kind: The kind of activity (Client, Server, Internal, etc.)
+            kind: The kind of activity. Accepts a string (e.g. ``"Client"``,
+                ``"Server"``, ``"Internal"``) or an ``opentelemetry.trace.SpanKind``
+                enum value directly.
             operation_name: The name of the operation being traced
             activity_name: The name of the activity for display purposes
             agent_details: Optional agent details
@@ -128,16 +130,20 @@ def __init__(
         if self._is_telemetry_enabled():
             tracer = self._get_tracer()
 
-            # Map string kind to SpanKind enum
-            activity_kind = SpanKind.INTERNAL
-            if kind.lower() == "client":
-                activity_kind = SpanKind.CLIENT
-            elif kind.lower() == "server":
-                activity_kind = SpanKind.SERVER
-            elif kind.lower() == "producer":
-                activity_kind = SpanKind.PRODUCER
-            elif kind.lower() == "consumer":
-                activity_kind = SpanKind.CONSUMER
+            # Resolve activity_kind from either a SpanKind enum or a string
+            if isinstance(kind, SpanKind):
+                activity_kind = kind
+            else:
+                # Map string kind to SpanKind enum
+                activity_kind = SpanKind.INTERNAL
+                if kind.lower() == "client":
+                    activity_kind = SpanKind.CLIENT
+                elif kind.lower() == "server":
+                    activity_kind = SpanKind.SERVER
+                elif kind.lower() == "producer":
+                    activity_kind = SpanKind.PRODUCER
+                elif kind.lower() == "consumer":
+                    activity_kind = SpanKind.CONSUMER
 
             # Get context for parent relationship
             # If parent_id is provided, parse it and use it as the parent context

tests/observability/core/test_execute_tool_scope.py

@@ -26,6 +26,7 @@
 from microsoft_agents_a365.observability.core.opentelemetry_scope import OpenTelemetryScope
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import SpanKind
 
 
 class TestExecuteToolScope(unittest.TestCase):
@@ -157,6 +158,26 @@ def test_execute_tool_scope_with_parent_id(self):
         span_parent_id = f"{span.parent.span_id:016x}"
         self.assertEqual(span_parent_id, parent_span_id)
 
+    def test_span_kind_defaults_to_internal(self):
+        """Test that ExecuteToolScope defaults to SpanKind.INTERNAL."""
+        scope = ExecuteToolScope.start(self.tool_details, self.agent_details, self.tenant_details)
+        scope.dispose()
+
+        finished_spans = self.span_exporter.get_finished_spans()
+        self.assertTrue(finished_spans, "Expected at least one span to be created")
+        self.assertEqual(finished_spans[-1].kind, SpanKind.INTERNAL)
+
+    def test_span_kind_override_to_client(self):
+        """Test that ExecuteToolScope accepts SpanKind.CLIENT override."""
+        scope = ExecuteToolScope.start(
+            self.tool_details, self.agent_details, self.tenant_details, span_kind=SpanKind.CLIENT
+        )
+        scope.dispose()
+
+        finished_spans = self.span_exporter.get_finished_spans()
+        self.assertTrue(finished_spans, "Expected at least one span to be created")
+        self.assertEqual(finished_spans[-1].kind, SpanKind.CLIENT)
+
 
 if __name__ == "__main__":
     # Run pytest only on the current file

tests/observability/core/test_invoke_agent_scope.py

@@ -30,6 +30,7 @@
 from microsoft_agents_a365.observability.core.opentelemetry_scope import OpenTelemetryScope
 from opentelemetry.sdk.trace.export import SimpleSpanProcessor
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+from opentelemetry.trace import SpanKind
 
 
 class TestInvokeAgentScope(unittest.TestCase):
@@ -199,6 +200,53 @@ def test_request_attributes_set_on_span(self):
                 input_messages,
             )
 
+    def test_invoke_agent_scope_span_kind(self):
+        """Test that InvokeAgentScope creates spans with the correct SpanKind."""
+        # Create scope
+        scope = InvokeAgentScope.start(
+            invoke_agent_details=self.invoke_details,
+            tenant_details=self.tenant_details,
+            request=self.test_request,
+        )
+
+        if scope is not None:
+            scope.dispose()
+
+        # Check that span was created with correct SpanKind
+        finished_spans = self.span_exporter.get_finished_spans()
+        self.assertTrue(finished_spans, "Expected at least one span to be created")
+
+        # Get the span and verify its kind
+        span = finished_spans[-1]
+        self.assertEqual(
+            span.kind,
+            SpanKind.CLIENT,
+            "InvokeAgentScope defaults to CLIENT spans (can be overridden with span_kind parameter)",
+        )
+
+        # Test SERVER span kind override
+        scope_server = InvokeAgentScope.start(
+            invoke_agent_details=self.invoke_details,
+            tenant_details=self.tenant_details,
+            request=self.test_request,
+            span_kind=SpanKind.SERVER,
+        )
+
+        if scope_server is not None:
+            scope_server.dispose()
+
+        # Check that SERVER span was created
+        finished_spans = self.span_exporter.get_finished_spans()
+        self.assertTrue(len(finished_spans) >= 2, "Expected at least two spans to be created")
+
+        # Get the most recent span and verify it's SERVER
+        server_span = finished_spans[-1]
+        self.assertEqual(
+            server_span.kind,
+            SpanKind.SERVER,
+            "InvokeAgentScope should create SERVER spans when span_kind parameter is set",
+        )
+
 
 if __name__ == "__main__":
     # Run pytest only on the current file