mseep-agentops 0.4.18__py3-none-any.whl

tests/unit/sdk/test_decorators.py

@@ -0,0 +1,763 @@
+from typing import AsyncGenerator
+import asyncio
+import pytest
+
+from agentops.sdk.decorators import agent, operation, session, workflow, task, tool
+from agentops.semconv import SpanKind
+from agentops.semconv.span_attributes import SpanAttributes
+from tests.unit.sdk.instrumentation_tester import InstrumentationTester
+from agentops.sdk.decorators.factory import create_entity_decorator
+
+
+class TestSpanNesting:
+    """Tests for proper nesting of spans in the tracing hierarchy."""
+
+    def test_operation_nests_under_agent(self, instrumentation: InstrumentationTester):
+        """Test that operation spans are properly nested under their agent spans."""
+
+        # Define the test agent with nested operations
+        @agent
+        class NestedAgent:
+            def __init__(self):
+                pass  # No logic needed
+
+            @operation
+            def nested_operation(self, message):
+                """Nested operation that should appear as a child of the agent"""
+                return f"Processed: {message}"
+
+            @operation
+            def main_operation(self):
+                """Main operation that calls the nested operation"""
+                # Call the nested operation
+                result = self.nested_operation("test message")
+                return result
+
+        # Test session with the agent
+        @session
+        def test_session():
+            agent = NestedAgent()
+            return agent.main_operation()
+
+        # Run the test with our instrumentor
+        result = test_session()
+
+        # Verify the result
+        assert result == "Processed: test message"
+
+        # Get all spans captured during the test
+        spans = instrumentation.get_finished_spans()
+
+        # Print detailed span information for debugging
+        print("\nDetailed span information:")
+        for i, span in enumerate(spans):
+            parent_id = span.parent.span_id if span.parent else "None"
+            span_id = span.context.span_id if span.context else "None"
+            print(f"Span {i}: name={span.name}, span_id={span_id}, parent_id={parent_id}")
+
+        # We should have 4 spans: session, agent, and two operations
+        assert len(spans) == 4
+
+        # Verify span kinds
+        session_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION
+        ]
+        agent_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT
+        ]
+        operation_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK
+        ]
+
+        assert len(session_spans) == 1
+        assert len(agent_spans) == 1
+        assert len(operation_spans) == 2
+
+        # Find the main_operation and nested_operation spans
+        main_operation = None
+        nested_operation = None
+
+        for span in operation_spans:
+            if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_operation":
+                main_operation = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_operation":
+                nested_operation = span
+
+        assert main_operation is not None, "main_operation span not found"
+        assert nested_operation is not None, "nested_operation span not found"
+
+        # Verify the session span is the root
+        session_span = session_spans[0]
+        assert session_span.parent is None
+
+        # Verify the agent span is a child of the session span
+        agent_span = agent_spans[0]
+        assert agent_span.parent is not None
+        assert session_span.context is not None
+        assert agent_span.parent.span_id == session_span.context.span_id
+
+        # Verify main_operation is a child of the agent span
+        assert main_operation.parent is not None
+        assert agent_span.context is not None
+        assert main_operation.parent.span_id == agent_span.context.span_id
+
+        # Verify nested_operation is a child of main_operation
+        assert nested_operation.parent is not None
+        assert main_operation.context is not None
+        assert nested_operation.parent.span_id == main_operation.context.span_id
+
+    def test_async_operations(self, instrumentation: InstrumentationTester):
+        """Test that async operations are properly nested."""
+
+        # Define the test agent with async operations
+        @agent
+        class AsyncAgent:
+            def __init__(self):
+                pass
+
+            @operation
+            async def nested_async_operation(self, message):
+                """Async operation that should appear as a child of the main operation"""
+                await asyncio.sleep(0.01)  # Small delay to simulate async work
+                return f"Processed async: {message}"
+
+            @operation
+            async def main_async_operation(self):
+                """Main async operation that calls the nested async operation"""
+                result = await self.nested_async_operation("test async")
+                return result
+
+        # Test session with the async agent
+        @session
+        async def test_async_session():
+            agent = AsyncAgent()
+            return await agent.main_async_operation()
+
+        # Run the async test
+        result = asyncio.run(test_async_session())
+
+        # Verify the result
+        assert result == "Processed async: test async"
+
+        # Get all spans captured during the test
+        spans = instrumentation.get_finished_spans()
+
+        # Print detailed span information for debugging
+        print("\nDetailed span information for async test:")
+        for i, span in enumerate(spans):
+            parent_id = span.parent.span_id if span.parent else "None"
+            span_id = span.context.span_id if span.context else "None"
+            print(f"Span {i}: name={span.name}, span_id={span_id}, parent_id={parent_id}")
+
+        # We should have 4 spans: session, agent, and two operations
+        assert len(spans) == 4
+
+        # Verify span kinds
+        session_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION
+        ]
+        agent_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT
+        ]
+        operation_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK
+        ]
+
+        assert len(session_spans) == 1
+        assert len(agent_spans) == 1
+        assert len(operation_spans) == 2
+
+        # Find the main_operation and nested_operation spans
+        main_operation = None
+        nested_operation = None
+
+        for span in operation_spans:
+            if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_async_operation":
+                main_operation = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_async_operation":
+                nested_operation = span
+
+        assert main_operation is not None, "main_async_operation span not found"
+        assert nested_operation is not None, "nested_async_operation span not found"
+
+        # Verify the session span is the root
+        session_span = session_spans[0]
+        assert session_span.parent is None
+
+        # Verify the agent span is a child of the session span
+        agent_span = agent_spans[0]
+        assert agent_span.parent is not None
+        assert session_span.context is not None
+        assert agent_span.parent.span_id == session_span.context.span_id
+
+        # Verify main_operation is a child of the agent span
+        assert main_operation.parent is not None
+        assert agent_span.context is not None
+        assert main_operation.parent.span_id == agent_span.context.span_id
+
+        # Verify nested_operation is a child of main_operation
+        assert nested_operation.parent is not None
+        assert main_operation.context is not None
+        assert nested_operation.parent.span_id == main_operation.context.span_id
+
+    def test_generator_operations(self, instrumentation: InstrumentationTester):
+        """Test that generator operations are properly nested."""
+
+        # Define the test agent with generator operations
+        @agent
+        class GeneratorAgent:
+            def __init__(self):
+                pass
+
+            @operation
+            def nested_generator(self, count):
+                """Generator operation that should appear as a child of the main operation"""
+                for i in range(count):
+                    yield f"Item {i}"
+
+            @operation
+            def main_generator_operation(self, count):
+                """Main operation that calls the nested generator"""
+                results = []
+                for item in self.nested_generator(count):
+                    results.append(item)
+                return results
+
+        # Test session with the generator agent
+        @session
+        def test_generator_session():
+            agent = GeneratorAgent()
+            return agent.main_generator_operation(3)
+
+        # Run the test
+        result = test_generator_session()
+
+        # Verify the result
+        assert result == ["Item 0", "Item 1", "Item 2"]
+
+        # Get all spans captured during the test
+        spans = instrumentation.get_finished_spans()
+
+        # Print detailed span information for debugging
+        print("\nDetailed span information for generator test:")
+        for i, span in enumerate(spans):
+            parent_id = span.parent.span_id if span.parent else "None"
+            span_id = span.context.span_id if span.context else "None"
+            print(f"Span {i}: name={span.name}, span_id={span_id}, parent_id={parent_id}")
+
+        # We should have 4 spans: session, agent, and two operations
+        assert len(spans) == 4
+
+        # Verify span kinds
+        session_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION
+        ]
+        agent_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT
+        ]
+        operation_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK
+        ]
+
+        assert len(session_spans) == 1
+        assert len(agent_spans) == 1
+        assert len(operation_spans) == 2
+
+        # Find the main_operation and nested_operation spans
+        main_operation = None
+        nested_operation = None
+
+        for span in operation_spans:
+            if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_generator_operation":
+                main_operation = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_generator":
+                nested_operation = span
+
+        assert main_operation is not None, "main_generator_operation span not found"
+        assert nested_operation is not None, "nested_generator span not found"
+
+        # Verify the session span is the root
+        session_span = session_spans[0]
+        assert session_span.parent is None
+
+        # Verify the agent span is a child of the session span
+        agent_span = agent_spans[0]
+        assert agent_span.parent is not None
+        assert session_span.context is not None
+        assert agent_span.parent.span_id == session_span.context.span_id
+
+        # Verify main_operation is a child of the agent span
+        assert main_operation.parent is not None
+        assert agent_span.context is not None
+        assert main_operation.parent.span_id == agent_span.context.span_id
+
+        # Verify nested_operation is a child of main_operation
+        assert nested_operation.parent is not None
+        assert main_operation.context is not None
+        assert nested_operation.parent.span_id == main_operation.context.span_id
+
+    def test_async_generator_operations(self, instrumentation: InstrumentationTester):
+        """Test that async generator operations are properly nested."""
+
+        # Define the test agent with async generator operations
+        @agent
+        class AsyncGeneratorAgent:
+            def __init__(self):
+                pass
+
+            @operation
+            async def nested_async_generator(self, count) -> AsyncGenerator[str, None]:
+                """Async generator operation that should appear as a child of the main operation"""
+                for i in range(count):
+                    await asyncio.sleep(0.01)  # Small delay to simulate async work
+                    yield f"Async Item {i}"
+
+            @operation
+            async def main_async_generator_operation(self, count):
+                """Main async operation that calls the nested async generator"""
+                results = []
+                async for item in self.nested_async_generator(count):
+                    results.append(item)
+                return results
+
+        # Test session with the async generator agent
+        @session
+        async def test_async_generator_session():
+            agent = AsyncGeneratorAgent()
+            return await agent.main_async_generator_operation(3)
+
+        # Run the async test
+        result = asyncio.run(test_async_generator_session())
+
+        # Verify the result
+        assert result == ["Async Item 0", "Async Item 1", "Async Item 2"]
+
+        # Get all spans captured during the test
+        spans = instrumentation.get_finished_spans()
+
+        # Print detailed span information for debugging
+        print("\nDetailed span information for async generator test:")
+        for i, span in enumerate(spans):
+            parent_id = span.parent.span_id if span.parent else "None"
+            span_id = span.context.span_id if span.context else "None"
+            print(f"Span {i}: name={span.name}, span_id={span_id}, parent_id={parent_id}")
+
+        # We should have 4 spans: session, agent, and two operations
+        assert len(spans) == 4
+
+        # Verify span kinds
+        session_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION
+        ]
+        agent_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT
+        ]
+        operation_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK
+        ]
+
+        assert len(session_spans) == 1
+        assert len(agent_spans) == 1
+        assert len(operation_spans) == 2
+
+        # Find the main_operation and nested_operation spans
+        main_operation = None
+        nested_operation = None
+
+        for span in operation_spans:
+            if (
+                span.attributes
+                and span.attributes.get(SpanAttributes.OPERATION_NAME) == "main_async_generator_operation"
+            ):
+                main_operation = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "nested_async_generator":
+                nested_operation = span
+
+        assert main_operation is not None, "main_async_generator_operation span not found"
+        assert nested_operation is not None, "nested_async_generator span not found"
+
+        # Verify the session span is the root
+        session_span = session_spans[0]
+        assert session_span.parent is None
+
+        # Verify the agent span is a child of the session span
+        agent_span = agent_spans[0]
+        assert agent_span.parent is not None
+        assert session_span.context is not None
+        assert agent_span.parent.span_id == session_span.context.span_id
+
+        # Verify main_operation is a child of the agent span
+        assert main_operation.parent is not None
+        assert agent_span.context is not None
+        assert main_operation.parent.span_id == agent_span.context.span_id
+
+        # Verify nested_operation is a child of main_operation
+        assert nested_operation.parent is not None
+        assert main_operation.context is not None
+        assert nested_operation.parent.span_id == main_operation.context.span_id
+
+    def test_complex_nesting(self, instrumentation: InstrumentationTester):
+        """Test complex nesting with multiple levels of operations."""
+
+        # Define the test agent with complex nesting
+        @agent
+        class ComplexAgent:
+            def __init__(self):
+                pass
+
+            @operation
+            def level3_operation(self, message):
+                """Level 3 operation (deepest)"""
+                return f"Level 3: {message}"
+
+            @operation
+            def level2_operation(self, message):
+                """Level 2 operation that calls level 3"""
+                result = self.level3_operation(message)
+                return f"Level 2: {result}"
+
+            @operation
+            def level1_operation(self, message):
+                """Level 1 operation that calls level 2"""
+                result = self.level2_operation(message)
+                return f"Level 1: {result}"
+
+        # Test session with the complex agent
+        @session
+        def test_complex_session():
+            agent = ComplexAgent()
+            return agent.level1_operation("test message")
+
+        # Run the test
+        result = test_complex_session()
+
+        # Verify the result
+        assert result == "Level 1: Level 2: Level 3: test message"
+
+        # Get all spans captured during the test
+        spans = instrumentation.get_finished_spans()
+
+        # Print detailed span information for debugging
+        print("\nDetailed span information for complex nesting test:")
+        for i, span in enumerate(spans):
+            parent_id = span.parent.span_id if span.parent else "None"
+            span_id = span.context.span_id if span.context else "None"
+            print(f"Span {i}: name={span.name}, span_id={span_id}, parent_id={parent_id}")
+
+        # We should have 5 spans: session, agent, and three operations
+        assert len(spans) == 5
+
+        # Verify span kinds
+        session_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION
+        ]
+        agent_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.AGENT
+        ]
+        operation_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK
+        ]
+
+        assert len(session_spans) == 1
+        assert len(agent_spans) == 1
+        assert len(operation_spans) == 3
+
+        # Find the operation spans
+        level1_operation = None
+        level2_operation = None
+        level3_operation = None
+
+        for span in operation_spans:
+            if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "level1_operation":
+                level1_operation = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "level2_operation":
+                level2_operation = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "level3_operation":
+                level3_operation = span
+
+        assert level1_operation is not None, "level1_operation span not found"
+        assert level2_operation is not None, "level2_operation span not found"
+        assert level3_operation is not None, "level3_operation span not found"
+
+        # Verify the session span is the root
+        session_span = session_spans[0]
+        assert session_span.parent is None
+
+        # Verify the agent span is a child of the session span
+        agent_span = agent_spans[0]
+        assert agent_span.parent is not None
+        assert session_span.context is not None
+        assert agent_span.parent.span_id == session_span.context.span_id
+
+        # Verify level1_operation is a child of the agent span
+        assert level1_operation.parent is not None
+        assert agent_span.context is not None
+        assert level1_operation.parent.span_id == agent_span.context.span_id
+
+        # Verify level2_operation is a child of level1_operation
+        assert level2_operation.parent is not None
+        assert level1_operation.context is not None
+        assert level2_operation.parent.span_id == level1_operation.context.span_id
+
+        # Verify level3_operation is a child of level2_operation
+        assert level3_operation.parent is not None
+        assert level2_operation.context is not None
+        assert level3_operation.parent.span_id == level2_operation.context.span_id
+
+    def test_workflow_and_task_nesting(self, instrumentation: InstrumentationTester):
+        """Test that workflow and task decorators create proper span nesting."""
+
+        # Define a workflow with tasks
+        @workflow
+        def data_processing_workflow(data):
+            """Main workflow that processes data through multiple tasks"""
+            result = process_input(data)
+            result = transform_data(result)
+            return result
+
+        @task
+        def process_input(data):
+            """Task to process input data"""
+            return f"Processed: {data}"
+
+        @task
+        def transform_data(data):
+            """Task to transform processed data"""
+            return f"Transformed: {data}"
+
+        # Test session with the workflow
+        @session
+        def test_workflow_session():
+            return data_processing_workflow("test data")
+
+        # Run the test
+        result = test_workflow_session()
+
+        # Verify the result
+        assert result == "Transformed: Processed: test data"
+
+        # Get all spans captured during the test
+        spans = instrumentation.get_finished_spans()
+
+        # Print detailed span information for debugging
+        print("\nDetailed span information for workflow and task test:")
+        for i, span in enumerate(spans):
+            parent_id = span.parent.span_id if span.parent else "None"
+            span_id = span.context.span_id if span.context else "None"
+            print(f"Span {i}: name={span.name}, span_id={span_id}, parent_id={parent_id}")
+
+        # We should have 4 spans: session, workflow, and two tasks
+        assert len(spans) == 4
+
+        # Verify span kinds
+        session_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.SESSION
+        ]
+        workflow_spans = [
+            s
+            for s in spans
+            if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.WORKFLOW
+        ]
+        task_spans = [
+            s for s in spans if s.attributes and s.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TASK
+        ]
+
+        assert len(session_spans) == 1
+        assert len(workflow_spans) == 1
+        assert len(task_spans) == 2
+
+        # Find the workflow and task spans
+        workflow_span = None
+        process_task = None
+        transform_task = None
+
+        for span in spans:
+            if span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "data_processing_workflow":
+                workflow_span = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "process_input":
+                process_task = span
+            elif span.attributes and span.attributes.get(SpanAttributes.OPERATION_NAME) == "transform_data":
+                transform_task = span
+
+        assert workflow_span is not None, "workflow span not found"
+        assert process_task is not None, "process_input task span not found"
+        assert transform_task is not None, "transform_data task span not found"
+
+        # Verify the session span is the root
+        session_span = session_spans[0]
+        assert session_span.parent is None
+
+        # Verify the workflow span is a child of the session span
+        assert workflow_span.parent is not None
+        assert session_span.context is not None
+        assert workflow_span.parent.span_id == session_span.context.span_id
+
+        # Verify process_task is a child of the workflow span
+        assert process_task.parent is not None
+        assert workflow_span.context is not None
+        assert process_task.parent.span_id == workflow_span.context.span_id
+
+        # Verify transform_task is a child of the workflow span
+        assert transform_task.parent is not None
+        assert workflow_span.context is not None
+        assert transform_task.parent.span_id == workflow_span.context.span_id
+
+
+@pytest.mark.asyncio
+async def test_async_context_manager():
+    """
+    Tests async context manager functionality (__aenter__, __aexit__).
+    """
+
+    # Create a simple decorated class
+    @create_entity_decorator("test")
+    class TestClass:
+        def __init__(self):
+            self.value = 42
+
+    # Cover __aenter__ and __aexit__ (normal exit)
+    async with TestClass() as instance:
+        assert hasattr(instance, "_agentops_active_span")
+        assert instance._agentops_active_span is not None
+
+    # Cover __aenter__ and __aexit__ (exceptional exit)
+    with pytest.raises(ValueError):
+        async with TestClass() as instance:
+            raise ValueError("Trigger exception for __aexit__ coverage")
+
+
+class TestToolDecorator:
+    """Tests for the tool decorator functionality."""
+
+    @pytest.fixture
+    def agent_class(self):
+        @agent
+        class TestAgent:
+            @tool(cost=0.01)
+            def process_item(self, item):
+                return f"Processed {item}"
+
+            @tool(cost=0.02)
+            async def async_process_item(self, item):
+                await asyncio.sleep(0.1)
+                return f"Async processed {item}"
+
+            @tool(cost=0.03)
+            def generator_process_items(self, items):
+                for item in items:
+                    yield self.process_item(item)
+
+            @tool(cost=0.04)
+            async def async_generator_process_items(self, items):
+                for item in items:
+                    await asyncio.sleep(0.1)
+                    yield await self.async_process_item(item)
+
+        return TestAgent()
+
+    def test_sync_tool_cost(self, agent_class, instrumentation: InstrumentationTester):
+        """Test synchronous tool with cost attribute."""
+        result = agent_class.process_item("test")
+
+        assert result == "Processed test"
+
+        spans = instrumentation.get_finished_spans()
+        tool_span = next(
+            span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL
+        )
+        assert tool_span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.01
+
+    @pytest.mark.asyncio
+    async def test_async_tool_cost(self, agent_class, instrumentation: InstrumentationTester):
+        """Test asynchronous tool with cost attribute."""
+        result = await agent_class.async_process_item("test")
+
+        assert result == "Async processed test"
+
+        spans = instrumentation.get_finished_spans()
+        tool_span = next(
+            span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL
+        )
+        assert tool_span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02
+
+    def test_generator_tool_cost(self, agent_class, instrumentation: InstrumentationTester):
+        """Test generator tool with cost attribute."""
+        items = ["item1", "item2", "item3"]
+        results = list(agent_class.generator_process_items(items))
+
+        assert len(results) == 3
+        assert results[0] == "Processed item1"
+        assert results[1] == "Processed item2"
+        assert results[2] == "Processed item3"
+
+        spans = instrumentation.get_finished_spans()
+        tool_spans = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL]
+        assert len(tool_spans) == 4  # Only one span for the generator
+        assert tool_spans[0].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.01
+        assert tool_spans[3].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.03
+
+    @pytest.mark.asyncio
+    async def test_async_generator_tool_cost(self, agent_class, instrumentation: InstrumentationTester):
+        """Test async generator tool with cost attribute."""
+        items = ["item1", "item2", "item3"]
+        results = [result async for result in agent_class.async_generator_process_items(items)]
+
+        assert len(results) == 3
+        assert results[0] == "Async processed item1"
+        assert results[1] == "Async processed item2"
+        assert results[2] == "Async processed item3"
+
+        spans = instrumentation.get_finished_spans()
+        tool_span = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL]
+        assert len(tool_span) == 4  # Only one span for the generator
+        assert tool_span[0].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02
+        assert tool_span[3].attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.04
+
+    def test_multiple_tool_calls(self, agent_class, instrumentation: InstrumentationTester):
+        """Test multiple calls to the same tool."""
+        for i in range(3):
+            result = agent_class.process_item(f"item{i}")
+            assert result == f"Processed item{i}"
+
+        spans = instrumentation.get_finished_spans()
+        tool_spans = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL]
+        assert len(tool_spans) == 3
+        for span in tool_spans:
+            assert span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.01
+
+    @pytest.mark.asyncio
+    async def test_parallel_tool_calls(self, agent_class, instrumentation: InstrumentationTester):
+        """Test parallel execution of async tools."""
+        results = await asyncio.gather(
+            agent_class.async_process_item("item1"),
+            agent_class.async_process_item("item2"),
+            agent_class.async_process_item("item3"),
+        )
+
+        assert len(results) == 3
+        assert results[0] == "Async processed item1"
+        assert results[1] == "Async processed item2"
+        assert results[2] == "Async processed item3"
+
+        spans = instrumentation.get_finished_spans()
+        tool_spans = [span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL]
+        assert len(tool_spans) == 3
+        for span in tool_spans:
+            assert span.attributes.get(SpanAttributes.LLM_USAGE_TOOL_COST) == 0.02
+
+    def test_tool_without_cost(self, agent_class, instrumentation: InstrumentationTester):
+        """Test tool without cost parameter."""
+
+        @tool
+        def no_cost_tool(self):
+            return "No cost tool result"
+
+        result = no_cost_tool(agent_class)
+
+        assert result == "No cost tool result"
+
+        spans = instrumentation.get_finished_spans()
+        tool_span = next(
+            span for span in spans if span.attributes.get(SpanAttributes.AGENTOPS_SPAN_KIND) == SpanKind.TOOL
+        )
+        assert SpanAttributes.LLM_USAGE_TOOL_COST not in tool_span.attributes