KG-178. Update Open Telemetry feature

Author: sdubov

agents/agents-core/src/commonMain/kotlin/ai/koog/agents/core/agent/StatefulSingleUseAIAgent.kt

@@ -85,30 +85,32 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
         }
 
         val runId = Uuid.random().toString()
-        val eventId = Uuid.random().toString()
-        val context = prepareContext(agentInput, runId, eventId)
 
-        return context.withPreparedPipeline {
+        // Unique identifier for a group of agent-run events
+        val agentRunEventId = Uuid.random().toString()
+        val context = prepareContext(agentInput, runId, agentRunEventId)
+
+        return withPreparedPipeline(context, agentRunEventId) {
             agentStateMutex.withLock {
                 @OptIn(InternalAgentsApi::class)
                 state = State.Running(context.parentContext ?: context)
             }
 
             logger.debug { formatLog(id, runId, "Starting agent execution") }
-            pipeline.onAgentStarting<Input, Output>(eventId, context.executionInfo, runId, this@StatefulSingleUseAIAgent, context)
+            pipeline.onAgentStarting<Input, Output>(agentRunEventId, context.executionInfo, runId, this@StatefulSingleUseAIAgent, context)
 
             val result = try {
                 @Suppress("UNCHECKED_CAST")
                 strategy.execute(context = context, input = agentInput)
             } catch (e: Throwable) {
                 logger.error(e) { "Execution exception reported by server!" }
-                pipeline.onAgentExecutionFailed(eventId, context.executionInfo, id, runId, e)
+                pipeline.onAgentExecutionFailed(agentRunEventId, context.executionInfo, id, runId, e)
                 agentStateMutex.withLock { state = State.Failed(e) }
                 throw e
             }
 
             logger.debug { formatLog(id, runId, "Finished agent execution") }
-            pipeline.onAgentCompleted(eventId, context.executionInfo, id, runId, result)
+            pipeline.onAgentCompleted(agentRunEventId, context.executionInfo, id, runId, result)
 
             agentStateMutex.withLock {
                 state = if (result != null) {
@@ -170,33 +172,34 @@ public abstract class StatefulSingleUseAIAgent<Input, Output, TContext : AIAgent
     protected fun formatLog(agentId: String, runId: String, message: String): String =
         "[agent id: $agentId, run id: $runId] $message"
 
+    //region Private Methods
+
     /**
      * Executes the provided block within a prepared pipeline context.
      * Ensures that the necessary feature resources are initialized before the block is invoked
      * and properly cleaned up after finishing the block.
      *
      * @return The result of the block's execution.
      */
-    private suspend fun <T> AIAgentContext.withPreparedPipeline(block: suspend (eventId: String) -> T): T {
-        require(executionInfo.parent == null) {
+    private suspend fun <T> withPreparedPipeline(context: AIAgentContext, eventId: String, block: suspend () -> T): T {
+        require(context.executionInfo.parent == null) {
             "withPreparedPipeline() should be called from a top level agent context."
         }
 
-        // Unique id for a group of agent-related events
-        val eventId = Uuid.random().toString()
-
         return try {
             pipeline.prepareFeatures()
-            block.invoke(eventId)
+            block.invoke()
         } finally {
             pipeline.onAgentClosing(
                 eventId = eventId,
-                executionInfo = executionInfo.parent ?: executionInfo,
+                executionInfo = context.executionInfo.parent ?: context.executionInfo,
                 agentId = id
             )
             pipeline.closeAllFeaturesMessageProcessors()
         }
     }
+
+    //endregion Private Methods
 }
 
 /**

agents/agents-core/src/commonTest/kotlin/ai/koog/agents/core/feature/AIAgentPipelineTest.kt

@@ -777,6 +777,75 @@ class AIAgentPipelineTest {
         assertContentEquals(expectedEvents, actualEvents)
     }
 
+    @Test
+    fun `test AgentExecutionInfo path with loop nodes invocation`() = runTest {
+        val interceptedEvents = mutableListOf<String>()
+        val interceptedRunIds = mutableListOf<String>()
+
+        val agentId = "test-agent-id"
+        val agentInput = "test input"
+
+        val strategyName = "test-strategy"
+        val nodeRootName = "node-root"
+        val nodeRootOutput = "Root node output"
+        val nodeExecuteName = "node-execute"
+        val nodeExecuteOutput = "Execute output"
+
+        var isExecuted = false
+
+        val strategy = strategy(strategyName) {
+            val nodeRoot by node<String, String>(nodeRootName) { it }
+            val nodeExecute by node<String, String>(nodeExecuteName) {
+                isExecuted = true
+                it
+            }
+
+            edge(nodeStart forwardTo nodeRoot)
+            edge(nodeRoot forwardTo nodeExecute onCondition { !isExecuted })
+            edge(nodeRoot forwardTo nodeFinish onCondition { isExecuted } transformed { nodeRootOutput })
+            edge(nodeExecute forwardTo nodeRoot transformed { nodeExecuteOutput })
+        }
+
+        createAgent(
+            strategy = strategy,
+            promptExecutor = getMockExecutor { }
+        ) {
+            install(TestFeature) {
+                events = interceptedEvents
+                runIds = interceptedRunIds
+            }
+        }.use { agent ->
+            agent.run(agentInput)
+        }
+
+        val actualEvents = interceptedEvents.filter { collectedEvent ->
+            collectedEvent.startsWith(NodeExecutionStarting::class.simpleName.toString()) ||
+                collectedEvent.startsWith(NodeExecutionCompleted::class.simpleName.toString()) ||
+                collectedEvent.startsWith(NodeExecutionFailed::class.simpleName.toString())
+        }
+
+        val expectedEvents = listOf(
+            "${NodeExecutionStarting::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, START_NODE_PREFIX)}, name: $START_NODE_PREFIX, input: $agentInput)",
+            "${NodeExecutionCompleted::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, START_NODE_PREFIX)}, name: $START_NODE_PREFIX, input: $agentInput, output: $agentInput)",
+            "${NodeExecutionStarting::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, nodeRootName)}, name: $nodeRootName, input: $agentInput)",
+            "${NodeExecutionCompleted::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, nodeRootName)}, name: $nodeRootName, input: $agentInput, output: $agentInput)",
+            "${NodeExecutionStarting::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, nodeExecuteName)}, name: $nodeExecuteName, input: $agentInput)",
+            "${NodeExecutionCompleted::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, nodeExecuteName)}, name: $nodeExecuteName, input: $agentInput, output: $agentInput)",
+            "${NodeExecutionStarting::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, nodeRootName)}, name: $nodeRootName, input: $nodeExecuteOutput)",
+            "${NodeExecutionCompleted::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, nodeRootName)}, name: $nodeRootName, input: $nodeExecuteOutput, output: $nodeExecuteOutput)",
+            "${NodeExecutionStarting::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, FINISH_NODE_PREFIX)}, name: $FINISH_NODE_PREFIX, input: $nodeRootOutput)",
+            "${NodeExecutionCompleted::class.simpleName} (path: ${agentExecutionPath(agentId, strategyName, FINISH_NODE_PREFIX)}, name: $FINISH_NODE_PREFIX, input: $nodeRootOutput, output: $nodeRootOutput)",
+        )
+
+        assertEquals(
+            expectedEvents.size,
+            actualEvents.size,
+            "Miss intercepted node events. Expected ${expectedEvents.size}, but received: ${actualEvents.size}"
+        )
+
+        assertContentEquals(expectedEvents, actualEvents)
+    }
+
     //endregion Execution Info
 
     //region Private Methods

agents/agents-features/agents-features-opentelemetry/README.md

@@ -158,7 +158,7 @@ The OpenTelemetry feature automatically adds various attributes to spans followi
 Some custom attributes specific to Koog are also added:
 
 - `koog.agent.strategy.name`: the name of the agent strategy
-- `koog.node.name`: the name of the node being executed
+- `koog.node.id`: the name of the node being executed
 
 ### Span Types
 
@@ -193,7 +193,7 @@ The OpenTelemetry feature creates different types of spans for various operation
     - Useful for: Understanding strategy flow and attributing LLM or tool work to a specific node.
     - Key attributes:
         - `gen_ai.conversation.id`
-        - `koog.node.name`
+        - `koog.node.id`
 
 4. **Inference Span**
     - Purpose: A single LLM call (prompt execution).